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

Question

我经常发现自己正在对数据帧进行分组应用，然后将结果与原始数据帧合并。这是一个例子。假设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