熊猫:将groupby-apply与join/merge结合起来

Question

我经常发现自己在数据帧上进行groupby-apply,然后将结果与原始数据帧合并.这是一个例子.假设df有A列和B列.我想添加另一列,其值是所有行的列B的总和,具有与当前行相同的列A值.以下是这项工作,但显然是次优的

df.join(df.groupby('A')['B'].sum(), on='A', rsuffix='_sum')

是否可以将原始索引保持在groupby-sum中？

Answer 1

你可以使用groupby.transform:

df['B_sum'] = df.groupby('A').B.transform('sum')

演示:

df = pd.DataFrame({
        'A': [1,1,2,2],
        'B': [1,2,3,4]
    })

df['B_sum'] = df.groupby('A').B.transform('sum')

df
#   A   B   B_sum
#0  1   1   3
#1  1   2   3
#2  2   3   7
#3  2   4   7