pandas DataFrame如何替代慢速apply实现累计余额计算

你当前使用的apply(axis=1)方案性能差的核心原因是逐行全表扫描，每处理一行都要遍历整个DataFrame做布尔筛选，时间复杂度为O(n²)，数据量达到十万级以上就会出现明显卡顿，百万级基本无法正常使用。

最高效的实现方式是利用pandas底层C实现的分组、累计求和向量化接口，时间复杂度为O(n)，相比apply方案性能提升可达数百到数千倍，完全适配大数据量场景。

实现步骤 #

1. 前置排序：先按交易时间升序排列数据，避免原始数据乱序导致计算错误
2. 分组累计：按币种分组，对交易金额做累计求和
3. 组内偏移：将累计和结果在组内向下偏移1位，首位空值填充为0，得到的就是不包含当前交易的历史金额总和

完整实现代码 #

import pandas as pd

# 原始数据构造
data = [[1653663281618, -583.8686, 'USD'],
       [1653741652125, -84.0381, 'USD'],
       [1653776860252, -33.8723, 'CHF'],
       [1653845294504, -465.4614, 'CHF'],
       [1653847155140, 22.285, 'USD'],
       [1653993629537, -358.04640000000006, 'USD']]
df = pd.DataFrame(data = data, columns = ['time', 'qty', 'currency_1'])

# 高性能计算balance列
df = df.sort_values('time', ascending=True).reset_index(drop=True)
df['balance'] = df.groupby('currency_1')['qty'].cumsum().groupby(df['currency_1']).shift(1, fill_value=0)

结果验证 #

运行后得到的计算结果和原apply方案完全一致，最终DataFrame内容如下：

注意事项 #

• 排序步骤不要省略，除非你100%确定原始数据已经严格按照time字段从小到大排列
• 偏移操作必须在分组内执行，不能直接对全局累计和做shift，否则会混入其他币种的交易数据
• 该方案已经是pandas生态下的最优实现，不需要额外用np.where手写逻辑，原生groupby+cumsum的底层优化程度远高于自定义numpy逻辑
• 如果数据量达到千万级以上，可以直接换用polars库，语法逻辑基本一致，内存占用和运行速度还能再提升3~10倍

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