基于两列唯一值创建新列：Pandas按组计算Ca-Cb

需求说明 #

现有包含A、B、C列的DataFrame，B列唯一值为'a'、'b'、'c'。需要创建新列D：按A列的每个分组，计算该组内B='a'对应的C值（记为Ca）减去B='b'对应的C值（记为Cb），并将计算得到的差值填充至该分组的所有行。

示例数据 #

import pandas as pd
data = {'A': ['X', 'X', 'X', 'Y', 'Y', 'Y', 'Z', 'Z', 'Z'], 
        'B': ['a', 'b', 'c', 'a', 'b', 'c', 'a', 'b' , 'c'],
        'C': [10,   20,  7,   30,  25,  7,   60,  50,   7],
        }
df = pd.DataFrame(data)

实现方法 #

方法一：使用groupby.transform直接计算 #

利用分组转换的特性，在每个分组内提取目标值计算差值，自动填充整组：

df['D'] = df.groupby('A').transform(
    lambda group: group[group['B'] == 'a']['C'].iloc[0] - group[group['B'] == 'b']['C'].iloc[0]
)

方法二：透视表计算差值后合并 #

先通过透视表生成每个分组的差值，再合并回原DataFrame：

# 生成每个A分组的Ca - Cb差值
diff_table = df.pivot(index='A', columns='B', values='C').eval('a - b').reset_index(name='D')
# 合并到原表
df = df.merge(diff_table, on='A', how='left')

结果验证 #

执行代码后，得到的DataFrame与预期一致：

# 输出结果
print(df)
"""
   A  B   C   D
0  X  a  10 -10
1  X  b  20 -10
2  X  c   7 -10
3  Y  a  30  -5
4  Y  b  25  -5
5  Y  c   7  -5
6  Z  a  60  10
7  Z  b  50  10
8  Z  c   7  10
"""

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