python - 选择性分组依据，np.where Pandas/Python 查询

Question

df_有

A  B C
1  1  10
1  1  10
1  2  5
1  3  6
1  4  7

df_想要

A  B C    D
1  1  10  20
1  1  10  20
1  2  5   5
1  3  6   6
1  4  7   7

仅当 B 列为 1 或 2 时，才尝试按 A 列和 B 列分组。如果 B 为 1 或 2，则将 D 列输出为 C 列的 SUM。否则保持设置列 D = 列 C。

示例代码报错:

df_want['D']=np.where((df_want['B'].isin([1,2]), 
             df_want['A','B'].map(df_want.groupby(['A','B'])['C'].sum()), 
             df_want['C'])

问题出现在这里:df_want['A','B'].map.. 如果我只放 df_want.A.map 那么代码运行但输出错误。如果 col B 是 1 或 2，我只需要它来映射总和

Answer 1

将 sum 更改为 transform('sum')

np.where(df_want['B'].isin([1,2]), 
             df_want.groupby(['A','B'])['C'].transform('sum'), df_want['C'])