熊猫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进行处理，然后使用熊猫进行合并，但是我想知道是否有更快的方法？在SQL中，可以使用partition来实现，而在R中，我可以使用dplyr和mutate。熊猫里有类似的东西吗？

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