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 中可以使用 with 来实现partition，在 RI 中可以使用dplyrand 来实现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