如何对dataframe指定列做[-3,3]上下限约束的cumsum累加计算

首先注意一个常见误区：不要直接对全列cumsum()结果调用clip()，普通累加会记录超出[-3,3]区间的数值，后续计算会基于溢出的错误值，和逐次钳制的预期结果不符。

以下是两种可直接落地的实现：

方法1：纯pandas实现（适合小数据量） #

无需额外依赖，直接用expanding.apply逐窗口计算：

import pandas as pd

# 构造示例数据
df = pd.DataFrame({'val': [1, 1, 1, 1, 0, -1]})

# 定义逐次累加+钳制的逻辑
def bounded_cum(series):
    current = 0
    for num in series:
        current += num
        current = max(min(current, 3), -3)
    return current

# 生成结果列
df['bounded_cumsum'] = df['val'].expanding().apply(bounded_cum).astype(int)

运行后输出的bounded_cumsum列就是你预期的[1, 2, 3, 3, 3, 2]。

方法2：Numba加速实现（适合大数据量） #

如果数据量超过1万行，上面的循环实现性能较低，可以用Numba JIT编译优化速度：

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

# 编译加速的累加函数
@jit(nopython=True)
def bounded_cumsum_func(arr, min_limit=-3, max_limit=3):
    res_arr = np.empty_like(arr)
    current_sum = 0
    for idx in range(len(arr)):
        current_sum += arr[idx]
        if current_sum > max_limit:
            current_sum = max_limit
        elif current_sum < min_limit:
            current_sum = min_limit
        res_arr[idx] = current_sum
    return res_arr

# 调用生成结果
df['bounded_cumsum'] = bounded_cumsum_func(df['val'].values)

该方法性能比纯pandas实现高百倍以上，适合大规模数据集使用。

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