如何无需逐行迭代实现Pandas多列分类的累计求和

现有如下Pandas DataFrame：

import pandas
import numpy
df = pandas.DataFrame( data= {'s1' :numpy.random.choice( ['A', 'B', 'C', 'D', 'E'], size=20 ),
                              's2' :numpy.random.choice( ['A', 'B', 'C', 'D', 'E'], size=20 ),
                              'val':numpy.random.randint(low=-1, high=3, size=20)},   )

需求：

• 生成ans1和ans2两列，分别基于s1、s2列的分类对val列做累计求和
• 首次出现的分类对应值为0，后续出现时取该分类此前所有val的累计和
• 分类可同时出现在s1或s2中（同一分类在两列出现时，累计值全局共享）

目前已通过逐行迭代+字典记录累计值实现，但效率较低。单列场景可使用groupby().cumsum().shift(1)实现，现寻求多列场景的通用向量化解决方案。

核心思路 #

把多列的分类统一转成长格式，用groupby全局计算每个分类的累计值，再映射回原表的对应列，全程用Pandas内置的向量化操作替代循环。

代码实现 #

import pandas as pd
import numpy as np

# 固定随机种子方便复现结果
np.random.seed(42)
df = pd.DataFrame( data= {'s1' :np.random.choice( ['A', 'B', 'C', 'D', 'E'], size=20 ),
                          's2' :np.random.choice( ['A', 'B', 'C', 'D', 'E'], size=20 ),
                          'val':np.random.randint(low=-1, high=3, size=20)},   )

# 1. 将宽表转长表，记录每行s1、s2的分类及其对应的val
long_df = pd.melt(df.reset_index(), 
                  id_vars=['index', 'val'], 
                  value_vars=['s1', 's2'], 
                  var_name='col', 
                  value_name='category')

# 2. 全局计算每个分类到当前行的累计和（不含当前行），首次出现填充0
long_df['cum_val'] = long_df.groupby('category')['val'].cumsum().shift(1).fillna(0)

# 3. 将计算好的累计值映射回原表的ans1、ans2列
ans1 = long_df[long_df['col'] == 's1'].set_index('index')['cum_val']
ans2 = long_df[long_df['col'] == 's2'].set_index('index')['cum_val']
df['ans1'] = ans1
df['ans2'] = ans2

效果说明 #

这个方案利用melt统一处理多列分类，通过groupby.cumsum实现全局累计，完全规避了逐行循环，效率比字典迭代提升几个数量级，且支持任意多分类列扩展。

内容的提问来源于stack exchange，提问作者Mudjud