Pandas添加dropped列：提取上一行code中当前行缺失的元素

问题背景 #

给定如下Pandas DataFrame：

import pandas as pd

data = {
    'accuracy': [0.773, 0.841, 0.862, 0.874, 0.883, 0.913],
    'code': [('D',),('D', 'F'),('B', 'D', 'F'),
             ('B', 'F', 'K'), ('B', 'F', 'I', 'K'),
             ('F', 'I', 'K')]
}
df = pd.DataFrame(data)

初始输出：

accuracy         code
0   0.773           (D,)
1   0.841         (D, F)
2   0.862      (B, D, F)
3   0.874      (B, F, K)
4   0.883   (B, F, I, K)
5   0.913      (F, I, K)

需要新增一列dropped，值为上一行code中存在但当前行code中不存在的元素，无符合元素时填'-'，预期结果：

accuracy        code    dropped
0   0.773           (D,)      -
1   0.841         (D, F)      -
2   0.862      (B, D, F)      -
3   0.874      (B, F, K)      D
4   0.883   (B, F, I, K)      -
5   0.913      (F, I, K)      B

实现代码 #

# 将code列转为集合，方便计算差集
df['code_set'] = df['code'].apply(set)
# 获取上一行的code集合
df['prev_code'] = df['code_set'].shift(1)

# 定义函数计算每一行的dropped值
def calculate_dropped(row):
    if pd.isna(row['prev_code']):
        return '-'
    # 计算上一行集合与当前行集合的差集
    dropped_elements = row['prev_code'] - row['code_set']
    # 差集为空则返回'-'，否则取出唯一元素
    return dropped_elements.pop() if dropped_elements else '-'

# 生成dropped列
df['dropped'] = df.apply(calculate_dropped, axis=1)
# 删除辅助列
df = df.drop(['code_set', 'prev_code'], axis=1)

# 输出结果
print(df)

代码说明 #

• df['code'].apply(set)：把每个元组转为集合类型，集合的差集操作能快速找出上一行有但当前行没有的元素
• df['code_set'].shift(1)：将code集合列向下偏移一行，得到每行对应的上一行集合数据，第一行会生成NaN
• calculate_dropped函数：处理边界情况（第一行无上行数据返回'-'），计算差集后，根据差集是否为空返回对应值
• 最后移除辅助列，得到目标格式的DataFrame

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