Присвоить среднее значение groupby столбцу при условии в пандах Python

Question

Привет, предположим, что у меня df_test, как показано ниже:

df_test = pd.DataFrame(dict(A=['bond1','bond1', 'bond2', 'bond2', 'bond3'],
 B=[-1, 1, 2, 3, 4]))

Я хочу создать столбец C, в котором, если B <0, это среднее значение группы B для группы A, если B> 0, это B. Метод, который я могу придумать:

group_mean = df_test.groupby('A')['B'].mean().reset_index()
df_merge = df_test.merge(group_mean, on=['A'], how='left')
df_merge['C'] = df_merge.apply(lambda x: x['B_y'] if x['B_x'] <0 else x['B_x'], axis=1)

Хотите знать, если это возможно сделать в одну строку, что-то вроде:

df_test['C'] = df_test.groupby('B').transform(...)

Спасибо за вашу помощь

Answer 1

IIUC:

df_test['C'] = df_test.groupby('A')['B'].transform('mean').mask(df['B']>0, df_test['B'])

Выход:

       A  B    C
0  bond1 -1  0.0
1  bond1  1  1.0
2  bond2  2  2.5
3  bond2  3  2.5
4  bond3  4  4.0