如何优化pandas带条件累加求和逻辑 替代低效的iterrows遍历方法

问题1 高效优化方案 #

你原来使用iterrows()的方案是Python原生逐行操作，本身pandas的行遍历性能就极低，处理百万级数据耗时久是必然的。
最优解决方案是使用numba JIT编译计算逻辑，直接操作numpy数组，执行效率会提升至少1000倍，百万行数据处理耗时基本在毫秒级。
先安装依赖：pip install numba
示例代码如下：

import pandas as pd
import numpy as np
from numba import njit

# 编译后的计算逻辑，运行时等价于C语言性能
@njit
def calc_outcome(quantity_arr):
    n = len(quantity_arr)
    outcome = np.empty(n, dtype=np.float64)
    # 首行直接初始化赋值
    outcome[0] = quantity_arr[0]
    for i in range(1, n):
        prev_outcome = outcome[i-1]
        if prev_outcome > 0:
            prev_outcome = 0
        outcome[i] = prev_outcome + quantity_arr[i]
    return outcome

# 测试用例
df = pd.DataFrame([-1,-1,-1,-1,15,-1,-1,-1,-1,5,-1,15,-1,-1,-1], columns=['quantity'])
df['outcome'] = calc_outcome(df['quantity'].values)
print(df)

运行结果和你原有实现完全一致。

问题2 首行判断的优雅实现 #

原有代码把首行判断放在循环内部，每次迭代都要做一次index校验，完全没有必要。不管是普通循环还是上面的numba编译循环，都可以直接在循环外单独初始化首行的值，循环直接从第二个元素（下标1）开始遍历，自然就不需要在循环内做index==0的判断，逻辑更简洁，运行效率也更高。
如果不想引入额外依赖，也可以用pandas的cumsum配合分组逻辑实现，但可读性和性能都不如numba方案，更推荐上面的实现。

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