Python Pandas DataFrame按分组基于列值计算自定义公式

你碰到的索引错位问题根源是筛选不同水果子集时，各子集保留的原始行索引没有对齐，没必要硬套eval写法，先把长表转宽表再做向量计算，全程自动对齐索引，代码最简洁不容易出错。

首先复现原始数据：

import pandas as pd
df = pd.DataFrame({
    'testid': (1,2,1,2,1,2),
    'Name': ('apple','apple','melon','melon','orange','orange'), 
    'A': (1,2,10,20,5,5), 
    'B': (1,5,4,2,3,1)
})

最优解法：透视后直接运算 #

用pivot把结构转换成以testid为索引、Name为列的宽表，每个单元格直接存储对应A、B列的数值，之后直接按公式计算即可，pandas会自动按testid索引对齐数据，不会出现错位：

# 生成透视表
pivot_df = df.pivot(index='testid', columns='Name', values=['A', 'B'])

# 按公式计算两个结果列
res = pd.DataFrame({
    'result_A': 2 * pivot_df[('A', 'apple')] - pivot_df[('A', 'melon')] - pivot_df[('A', 'orange')],
    'result_B': 2 * pivot_df[('B', 'apple')] - pivot_df[('B', 'melon')] - pivot_df[('B', 'orange')]
}).reset_index()

运行得到的结果完全匹配预期：

之前写法的问题说明 #

• 调用df.eval()时，默认的执行上下文是当前DataFrame的列名，直接写df[筛选条件]相当于查找名为df的列，自然会触发UndefinedVariableError
• 直接筛选不同Name的子集做算术运算时，apple子集的行索引是0、1，melon子集是2、3，orange子集是4、5，pandas做运算默认按行索引对齐，索引不匹配的位置会返回空值，根本算不出正确结果。

备选：分组apply写法 #

如果后续计算逻辑更复杂，也可以用groupby配合自定义函数实现，只是性能比透视法稍差：

def calc(group):
    apple = group.loc[group['Name']=='apple', ['A','B']].iloc[0]
    melon = group.loc[group['Name']=='melon', ['A','B']].iloc[0]
    orange = group.loc[group['Name']=='orange', ['A','B']].iloc[0]
    cal_res = 2*apple - melon - orange
    return pd.Series({'result_A': cal_res['A'], 'result_B': cal_res['B']})

res = df.groupby('testid').apply(calc).reset_index()

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