如何为二维数组按元素级不同窗口快速计算滚动和（支持前后方向）

我需要对二维数组按每行对应的不同窗口大小计算滚动和，支持向前或向后方向。目前实现的函数速度过慢（需调用数百甚至数千次），以下是当前代码：

import numpy as np

def rolling_sum(ar, window, direction="forward"):
    ar_sum = ar.copy().astype(float) # 默认初始窗口为1
    for w in np.unique(window):
        if w > 1: # 窗口为1时无需额外计算
            mask = (window == w) # 匹配当前窗口的行
            if direction=="forward":
                ar_sum[mask,:-w+1] += sum([ar[mask,col:ar.shape[1]+1-w+col] for col in range(1, w)])
                ar_sum[mask,-w+1:] = np.nan # 末尾无效区域设为NaN
            elif direction=="backward":
                ar_sum[mask,w-1:] += sum([ar[mask,w-1-col:ar.shape[1]-col] for col in range(1, w)])
                ar_sum[mask,:w-1] = np.nan # 开头无效区域设为NaN
    # 处理窗口为0的情况
    mask = window == 0
    ar_sum[mask] = 0
    return ar_sum

测试代码：

ar = np.random.randint(-100, 100, size=(10000, 50))
window = np.random.randint(0, 10, size=(10000))
rolling_sum(ar, window, direction="forward")

优化方案：利用前缀和实现向量化计算 #

原函数的核心瓶颈在于循环切片求和和逐窗口遍历的低效操作，改用前缀和可以将时间复杂度从O(NWC)降至O(N*C)（N为行数，W为窗口数，C为列数），大幅提升速度。

优化后代码 #

import numpy as np

def fast_rolling_sum(ar, window, direction="forward"):
    ar = ar.astype(float)
    n_rows, n_cols = ar.shape
    ar_sum = ar.copy()
    
    # 处理窗口为0的行
    zero_mask = window == 0
    ar_sum[zero_mask] = 0
    
    # 处理窗口>1的行
    valid_windows = np.unique(window[window > 1])
    for w in valid_windows:
        mask = window == w
        if not np.any(mask):
            continue
        
        if direction == "forward":
            # 计算前缀和，快速推导任意窗口的滚动和
            prefix = np.cumsum(ar[mask], axis=1)
            # 滚动和 = 前缀和[w-1:] - 前缀和[:-1]（补零处理首项）
            ar_sum[mask, :n_cols - w + 1] = prefix[:, w-1:] - np.pad(prefix[:, :-1], ((0,0), (1,0)), mode='constant')
            # 末尾w-1列无有效窗口，设为NaN
            ar_sum[mask, n_cols - w + 1:] = np.nan
        elif direction == "backward":
            # 反向计算前缀和，适配向后窗口
            reversed_ar = ar[mask, ::-1]
            prefix = np.cumsum(reversed_ar, axis=1)
            # 转换回原数组的有效列范围
            ar_sum[mask, w-1:] = prefix[:, w-1:] - np.pad(prefix[:, :-1], ((0,0), (1,0)), mode='constant')
            # 开头w-1列无有效窗口，设为NaN
            ar_sum[mask, :w-1] = np.nan
    
    return ar_sum

优化说明 #

1. 前缀和的优势：通过np.cumsum一次计算得到所有前缀和，后续任意窗口的和都可以通过前缀和的差值直接推导，避免了多次切片求和的冗余计算。
2. 向量化操作：所有行级计算均为numpy向量化操作，比Python循环快几个数量级。
3. 内存效率：避免了原函数中生成多个中间切片数组，减少内存占用。

性能测试 #

import time

ar = np.random.randint(-100, 100, size=(10000, 50))
window = np.random.randint(0, 10, size=(10000))

# 原函数耗时
start = time.time()
res_old = rolling_sum(ar, window, direction="forward")
print(f"原函数耗时: {time.time() - start:.4f}秒")

# 优化后函数耗时
start = time.time()
res_new = fast_rolling_sum(ar, window, direction="forward")
print(f"优化后函数耗时: {time.time() - start:.4f}秒")

# 验证结果一致性
print(np.allclose(res_old, res_new, equal_nan=True))  # 输出True表示结果一致

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