Pandas - 根据其他列创建总列

Question

我正在尝试创建一个总列,它根据第三列对另一列中的数字求和.我可以通过使用.groupby(),但这创建一个截断列,而我想要一个相同长度的列.

我的代码:

df = pd.DataFrame({'a':[1,2,2,3,3,3], 'b':[1,2,3,4,5,6]})
df['total'] = df.groupby(['a']).sum().reset_index()['b']

我的结果:

   a  b  total
0  1  1    1.0
1  2  2    5.0
2  2  3   15.0
3  3  4    NaN
4  3  5    NaN
5  3  6    NaN

我想要的结果:

   a  b  total
0  1  1    1.0
1  2  2    5.0
2  2  3    5.0
3  3  4   15.0
4  3  5   15.0
5  3  6   15.0

...每'a'列的总和与另一列相同.

Answer 1

从pandas中的groupby操作返回总和只会产生一个列,只要索引中的唯一项数.使用transform生成与原始数据帧长度相同("like-indexed")的列,而不执行任何合并.

df['total'] = df.groupby('a')['b'].transform(sum)
>>> df
   a  b  total
0  1  1      1
1  2  2      5
2  2  3      5
3  3  4     15
4  3  5     15
5  3  6     15