Pandas按月分组受限累计求和：现有方法失效且效率低，求解决方案

1. 原始DataFrame #

import pandas as pd
import numpy as np

data = {'c1':[-1,-1,1,1,np.nan,1,1,1,1,1,np.nan,-1],\
        'c2':[1,1,1,-1,1,1,-1,-1,1,-1,1,np.nan]}

index = pd.date_range('2000-01-01','2000-03-20', freq='W')

df = pd.DataFrame(index=index, data=data)

输出结果：

c1   c2
2000-01-02 -1.0  1.0
2000-01-09 -1.0  1.0
2000-01-16  1.0  1.0
2000-01-23  1.0 -1.0
2000-01-30  NaN  1.0
2000-02-06  1.0  1.0
2000-02-13  1.0 -1.0
2000-02-20  1.0 -1.0
2000-02-27  1.0  1.0
2000-03-05  1.0 -1.0
2000-03-12  NaN  1.0
2000-03-19 -1.0  NaN

2. 普通按月累计求和结果 #

df2 = df.groupby(df.index.to_period('m')).cumsum()

输出结果：

c1   c2
2000-01-02 -1.0  1.0
2000-01-09 -2.0  2.0
2000-01-16 -1.0  3.0
2000-01-23  0.0  2.0
2000-01-30  NaN  3.0
2000-02-06  1.0  1.0
2000-02-13  2.0  0.0
2000-02-20  3.0 -1.0
2000-02-27  4.0  0.0
2000-03-05  1.0 -1.0
2000-03-12  NaN  0.0
2000-03-19  0.0  NaN

3. 需求：受限累计求和 #

需要实现受限累计求和：当累计值超过3或低于0时不再继续增减，自定义函数如下：

def cumsum2(arr, low=-float('Inf'), high=float('Inf')):
    arr2 = np.copy(arr)
    sm = 0
    for index, elem in np.ndenumerate(arr):
        if not np.isnan(elem):
            sm += elem
            if sm > high:
                sm = high
            if sm < low:
                sm = low
        arr2[index] = sm
    return arr2

期望得到的结果：

c1   c2
2000-01-02  0.0  1.0
2000-01-09  0.0  2.0
2000-01-16  1.0  3.0
2000-01-23  2.0  2.0
2000-01-30  2.0  3.0
2000-02-06  1.0  1.0
2000-02-13  2.0  0.0
2000-02-20  3.0  0.0
2000-02-27  3.0  1.0
2000-03-05  1.0  0.0
2000-03-12  1.0  1.0
2000-03-19  0.0  1.0

4. 遇到的问题 #

尝试使用groupby+apply调用自定义函数，但无法正常运行，且大数据量下速度极慢：

df.groupby(df.index.to_period('m')).apply(lambda x: cumsum2(x, 0, 3))

问题分析与高效实现 #

问题出在哪里？ #

1. 参数结构不匹配：groupby.apply传入的是每组的DataFrame，而cumsum2是针对numpy一维数组设计的，直接传入二维DataFrame会导致np.ndenumerate遍历逻辑错误，返回结果无法匹配原DataFrame结构。
2. 性能瓶颈：apply本身是逐组循环，加上cumsum2里的Python级遍历，完全没有利用numpy/pandas的向量化运算优势，大数据量下效率极低。

高效实现方式 #

受限累计是状态依赖型运算（每一步结果依赖上一步），无法完全向量化，但可以用numba的JIT编译加速循环，同时结合groupby.transform保证结果结构匹配。

方案1：numba加速的最优实现

先安装numba：pip install numba，然后使用以下代码：

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

# 用numba编译一维数组的受限累计函数
@jit(nopython=True)
def bounded_cumsum(arr, low=0, high=3):
    result = np.empty_like(arr)
    current_sum = 0.0
    for i in range(len(arr)):
        val = arr[i]
        if not np.isnan(val):
            current_sum += val
            # 限制上下界
            current_sum = max(low, min(high, current_sum))
        result[i] = current_sum
    return result

# 按月份分组，对每列应用受限累计
df_result = df.groupby(df.index.to_period('m')).transform(lambda col: bounded_cumsum(col.values, 0, 3))

print(df_result)

运行结果完全符合期望，且大数据量下速度比纯Python循环快几十倍。

方案2：纯pandas实现（无numba）

如果无法安装numba，可以用expanding结合自定义逻辑，性能略逊于numba，但比原方案高效：

def bounded_cumsum_pd(col, low=0, high=3):
    def update(window):
        prev = window.iloc[-2] if len(window) > 1 else 0
        current = window.iloc[-1]
        if pd.isna(current):
            return prev
        new_val = prev + current
        return max(low, min(high, new_val))
    return col.expanding().apply(update, raw=False)

df_result = df.groupby(df.index.to_period('m')).apply(bounded_cumsum_pd, low=0, high=3)

方案优势 #

• numba JIT编译：将Python循环转换为机器码，执行速度接近C语言，彻底解决性能问题。
• transform正确应用：groupby.transform自动将每组的计算结果映射回原DataFrame的索引结构，避免了apply的结构匹配问题。

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