Pandas DataFrame基于前序ref1同ref2值生成新列的实现问题

实现代码 #

import pandas as pd
import numpy as np

# 构造原始DataFrame
df = pd.DataFrame({
    'ref1': [1,1,3,7,7],
    'ref2': [1,2,1,1,2],
    'value': [1,2,3,5,6],
})

# 步骤1：获取ref1按出现顺序的唯一列表，构造当前ref1到前一个ref1的映射关系
sorted_ref1 = df['ref1'].unique()
prev_ref1_map = {sorted_ref1[i]: sorted_ref1[i-1] for i in range(1, len(sorted_ref1))}

# 步骤2：构造每个ref1分组下，ref2到对应value的映射字典
ref2_value_map = df.groupby('ref1').apply(lambda g: g.set_index('ref2')['value'].to_dict()).to_dict()

# 步骤3：逐行匹配生成new_value列
df['new_value'] = df.apply(
    lambda row: ref2_value_map.get(prev_ref1_map.get(row['ref1']), {}).get(row['ref2'], np.nan),
    axis=1
)

输出结果 #

运行代码后得到的df如下：

ref1  ref2  value  new_value
0     1     1      1        NaN
1     1     2      2        NaN
2     3     1      3        1.0
3     7     1      5        3.0
4     7     2      6        NaN

逻辑说明 #

• 先按出现顺序提取所有唯一的ref1值，给每个ref1匹配它前一个出现的ref1值
• 提前把每个ref1分组下的ref2和对应value存储为字典，避免重复分组查询，提升运行效率
• 逐行匹配规则：先找当前行ref1对应的前一个ref1，再去前一个ref1的映射字典里找相同ref2对应的value，找不到匹配项就返回NaN

内容的提问来源于stack exchange，提问作者makpalan