在python pandas中的条件下为列分配groupby均值

Question

嗨，假设我的df_test如下：

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

我想创建一个列C，如果B <0，它是B groupby 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