如何用Pandas按列分组求和并将结果插入每行新列

问题 #

我用Python把Excel数据处理流程从几小时压缩到了几分钟，已经通过Pandas生成透视表，当前需要新增YearT列：按AY（年份）分组对RsvT列求和，将对应年份的总和填充到该年份的每一行，同时保留所有原有列。

现有代码：

import pandas as pd
import numpy as np
    
ld = pd.read_excel('LossData.xlsx')

ld['RsvT'] = ld['RsvLoss_ITD'] + ld['RsvExpense_ITD']
    
pivot = pd.pivot_table(ld, index = [ 'AY', 'Fusion_Program', 'Fusion_Office_Adj', 'AsLob', 'LosOccurredState'], values = ['RsvT'], aggfunc=np.sum)
pivot.to_excel("test.xlsx", merge_cells=False)
pt = pd.read_excel('test.xlsx')

当前输出（数据已模糊处理）：

AY  Fusion_Program  Fusion_Office_Adj  AsLob LosOccurredState  RsvT
0  2011            1999                 10     52               PA     5
1  2012            1100                 10     52               NJ     6
2  2013            1999                 10     52               FL     3
3  2014            1100                 10     52               PA     7
4  2014            1150                 10     52               FL     8
5  2014            1999                 10     52               FL     4
6  2014            1999                 10     52               PA     3
7  2014            5000                 10    171               IN     2

期望输出：

AY  Fusion_Program  Fusion_Office_Adj  AsLob LosOccurredState  RsvT  YearT
0  2011            1999                 10     52               PA     5      5
1  2012            1100                 10     52               NJ     6      6
2  2013            1999                 10     52               FL     3      3
3  2014            1100                 10     52               PA     7     24
4  2014            1150                 10     52               FL     8     24
5  2014            1999                 10     52               FL     4     24
6  2014            1999                 10     52               PA     3     24
7  2014            5000                 10    171               IN     2     24

YearT列是对应AY年份的RsvT总和，例如2014年总和为7+8+4+3+2=24，需将该值填充到2014年的每一行。已知可能需要用到groupby函数，但不清楚如何实现该格式的输出。

解决方案 #

不需要将透视表导出再重新读取，直接在内存中的透视表DataFrame上操作即可，用groupby结合transform方法可以快速实现需求，transform会自动将分组计算的结果匹配到原DataFrame的每一行，完美保留所有原有列。

优化后的代码：

import pandas as pd
import numpy as np
    
ld = pd.read_excel('LossData.xlsx')

# 计算RsvT列
ld['RsvT'] = ld['RsvLoss_ITD'] + ld['RsvExpense_ITD']

# 生成透视表并将索引列转为普通数据列
pivot = pd.pivot_table(
    ld, 
    index=['AY', 'Fusion_Program', 'Fusion_Office_Adj', 'AsLob', 'LosOccurredState'], 
    values=['RsvT'], 
    aggfunc=np.sum
).reset_index()

# 新增YearT列：按AY分组求和后填充到对应年份的每一行
pivot['YearT'] = pivot.groupby('AY')['RsvT'].transform('sum')

# 输出到Excel
pivot.to_excel("test.xlsx", merge_cells=False, index=False)

关键说明 #

• reset_index()：透视表默认将index参数中的列设为行索引，使用该方法可将其转回普通数据列，避免后续分组操作的索引问题，同时匹配期望的输出格式。
• groupby('AY')['RsvT'].transform('sum')：
  1. 按AY字段分组，计算每组RsvT的总和；
  2. transform方法会将分组得到的总和，自动映射到原DataFrame中对应AY的每一行，无需手动进行数据合并。
• 避免重复读写Excel：直接在内存中处理透视表，省去导出再读取的冗余步骤，进一步提升处理效率。

内容的提问来源于Stack Exchange，提问作者SaicheS