Pandas groupby，求和并填充原始数据框

Question

这是我原来的df

import pandas as pd
df_1 = pd.DataFrame({'color': ['blue', 'blue', 'yellow', 'yellow'], 'count': [1,3,4,5]})

color   count   
blue    1   
blue    3   
yellow  4   
yellow  5   

我想按颜色列和总和计数列分组，然后用结果填充原始数据框。所以最终结果应该是这样的：

df_2 = pd.DataFrame({'color': ['blue', 'blue', 'yellow', 'yellow'], 'count': [1,3,4,5], 
                     'total_per_color': [4,4,9,9]})


color   count   total_per_color
blue    1       4
blue    3       4
yellow  4       9
yellow  5       9

我可以使用 groupby 和 sum 来完成，然后使用 pandas 进行合并，但我想知道是否有更快的方法来做到这一点？在 SQL 中可以使用 partition 来实现，在 R 中我可以使用 dplyr 和 mutate。 pandas 有类似的东西吗？

Answer 1

将transform 与groupby 一起使用

df_1['total_per_color']=df_1.groupby('color')['count'].transform('sum')
df_1
Out[886]: 
    color  count  total_per_color
0    blue      1                4
1    blue      3                4
2  yellow      4                9
3  yellow      5                9