如何为Pandas DataFrame添加列提取tuple列中的新增元素

现有如下df：

import pandas as pd

df = pd.DataFrame(
    {'loss': [0.044, 0.044, 0.038, 0.037, 0.036], 
  'code': ["('ac',)", "('ac', 'be')", "('ab', 'ac', 'be')",
              "('ab', 'ac', 'be', 'fi')", "('ab', 'ac', 'be', 'de', 'fi')"]}
)

输出结果：

loss                            code
0   0.044                          ('ac',)
1   0.044                     ('ac', 'be')
2   0.038               ('ab', 'ac', 'be')
3   0.037         ('ab', 'ac', 'be', 'fi')
4   0.036   ('ab', 'ac', 'be', 'de', 'fi')

需要新增一列added-code，用来存储code列中每行相比前一行新增的元素，预期结果如下：

loss                            code   added-code
0   0.044                          ('ac',)     ac
1   0.044                     ('ac', 'be')     be
2   0.038               ('ab', 'ac', 'be')     ab
3   0.037         ('ab', 'ac', 'be', 'fi')     fi
4   0.036   ('ab', 'ac', 'be', 'de', 'fi')     de

解决方案 #

1. 解析code列为集合 #

code列的内容是字符串格式的元组，先把它们转换成Python集合，方便后续计算差异：

import ast

df['code_set'] = df['code'].apply(lambda x: set(ast.literal_eval(x)))

2. 计算每行与前一行的新增元素 #

用shift()获取前一行的集合，通过集合差集得到新增元素；第一行没有前一行，直接取它的唯一元素：

df['added-code'] = df.apply(
    lambda row: next(iter(row['code_set'])) if row.name == 0 
    else next(iter(row['code_set'] - df['code_set'].shift()[row.name])),
    axis=1
)

3. 移除临时列（可选） #

如果不需要code_set临时列，直接删除：

df.drop('code_set', axis=1, inplace=True)

完整代码 #

import pandas as pd
import ast

df = pd.DataFrame(
    {'loss': [0.044, 0.044, 0.038, 0.037, 0.036], 
  'code': ["('ac',)", "('ac', 'be')", "('ab', 'ac', 'be')",
              "('ab', 'ac', 'be', 'fi')", "('ab', 'ac', 'be', 'de', 'fi')"]}
)

# 解析code列为集合
df['code_set'] = df['code'].apply(lambda x: set(ast.literal_eval(x)))

# 计算新增元素
df['added-code'] = df.apply(
    lambda row: next(iter(row['code_set'])) if row.name == 0 
    else next(iter(row['code_set'] - df['code_set'].shift()[row.name])),
    axis=1
)

# 删除临时列
df.drop('code_set', axis=1, inplace=True)

print(df)

运行后就能得到预期的结果。

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