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, а затем объединить с помощью панд, но мне интересно, есть ли более быстрый способ сделать это?В SQL это можно сделать с помощью partition, в RI можно использовать 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