операции с Python над несколькими столбцами после группы

Question

У меня есть датафрейм, как показано ниже:

df = {'col_1': [1,2,3,4,5,6,7,8,9,10],
      'col_2': [1,1,1,0,0,1,0,0,1,1],
      'col_3':['A','A','A','A','A','B','B','B','B','B']}
df = pd.DataFrame(df)

Я хочу что-то вроде:

df['col_4'] =  df.groupby(['col_3'])['col_1','col_2'].transform((col_1*col_2)/sum(col_2))

, которое является умножением col_1 и col_2, затем делится на суммуcol_2 в каждой группе, определенной col_3.Ожидаемый результат должен быть:

df
   col_1  col_2 col_3  col_4
0      1      1     A   2.00
1      2      1     A   2.00
2      3      1     A   2.00
3      4      0     A   2.00
4      5      0     A   2.00
5      6      1     B   8.33
6      7      0     B   8.33
7      8      0     B   8.33
8      9      1     B   8.33
9     10      1     B   8.33

Кто-нибудь знает, как туда добраться?

Answer 1

Предполагая, что вы хотите рассчитать сумму из col_1 * col_2 в группе, вы можете определить GroupBy объект и использовать transform + sum дважды:

g = df.eval('prod=col_1*col_2').groupby('col_3')
df['col_4'] = g['prod'].transform('sum') / g['col_2'].transform('sum')

print(df)

   col_1  col_2 col_3     col_4
0      1      1     A  2.000000
1      2      1     A  2.000000
2      3      1     A  2.000000
3      4      0     A  2.000000
4      5      0     A  2.000000
5      6      1     B  8.333333
6      7      0     B  8.333333
7      8      0     B  8.333333
8      9      1     B  8.333333
9     10      1     B  8.333333