python - 对某些列进行 groupby 和 sum 与 pandas，同时还包括其他列

Question

我有以下数据:

   import pandas as pd
x4 = pd.DataFrame({"ID": [101,101, 102, 103, 104, 105],
                   "Prob": [1, 1,1, 1, 1, 1],
                   "Ef": [0,2, 0, 0, 0.25, 0.29],
                   "W": [2, 2,3, 4, 5, 6],
                   "EC": [0, 0,0, 0, 1.6, 2],
                   "Rand": [11, 12,12, 13, 14, 15]})

我想获取sum(Prob * Ef)按ID，然后仅保留列ID、包含 sum 的列、EC 列和 W 列。

所以最后我想要这个:

            ID  sum_column EC       W
1:          101 2.00       0.0      2
2:          101 2.00       0.0      2
3:          102 0.00       0.0      3
4:          103 0.00       0.0      4
5:          104 0.25       1.6      5
6:          105 0.29       2.0      6

我已经尝试过这个:x4.loc[:, ['EC','W','ID','Prob','Ef']].groupby('ID').sum(Prob *Ef)

但是不起作用

Answer 1

使用GroupBy.transform按列相乘:

x4['sum_column'] = x4['Prob'].mul(x4['Ef']).groupby(x4['ID']).transform('sum')
x4 = x4.drop(['Ef','Prob', 'Rand'], axis=1)
print (x4)
    ID  W   EC  sum_column
0  101  2  0.0        2.00
1  101  2  0.0        2.00
2  102  3  0.0        0.00
3  103  4  0.0        0.00
4  104  5  1.6        0.25
5  105  6  2.0        0.29

如果列的顺序很重要，请使用 insert :

x4.insert(1, 'sum_column',  x4['Prob'].mul(x4['Ef']).groupby(x4['ID']).transform('sum'))
x4 = x4.drop(['Ef','Prob', 'Rand'], axis=1)
print (x4)
    ID  sum_column  W   EC
0  101        2.00  2  0.0
1  101        2.00  2  0.0
2  102        0.00  3  0.0
3  103        0.00  4  0.0
4  104        0.25  5  1.6
5  105        0.29  6  2.0