如何使用Numpy更高效计算滚动加权求和？

需求说明 #

需要计算一种"滚动加权和"：对于数组A中的每个元素，根据给定的权重-区间规则（列表d），找到A中落在对应偏移区间内的元素，将其对应的B数组元素求和后乘以权重，最终累加所有规则的结果。现有实现内存占用过高，希望通过Numpy高级函数优化性能。

示例数据 #

import numpy as np

A = np.append(np.linspace(0, 1, 10), np.linspace(1.1, 2, 30))
np.random.seed(0)

B = np.random.randint(3, size=40) + 1

# 权重-区间规则列表：[(权重, (偏移下限, 偏移上限))]
d = [(1, (-0.25, -0.20)), (0.5, (-0.20, -0.10)), (2, (-0.10, 0.15))]

数组A #

array([0.        , 0.11111111, 0.22222222, 0.33333333, 0.44444444,
       0.55555556, 0.66666667, 0.77777778, 0.88888889, 1.        ,
       1.1       , 1.13103448, 1.16206897, 1.19310345, 1.22413793,
       1.25517241, 1.2862069 , 1.31724138, 1.34827586, 1.37931034,
       1.41034483, 1.44137931, 1.47241379, 1.50344828, 1.53448276,
       1.56551724, 1.59655172, 1.62758621, 1.65862069, 1.68965517,
       1.72068966, 1.75172414, 1.78275862, 1.8137931 , 1.84482759,
       1.87586207, 1.90689655, 1.93793103, 1.96896552, 2.        ])

数组B #

array([1, 2, 1, 2, 2, 3, 1, 3, 1, 1, 1, 3, 2, 3, 3, 1, 2, 2, 2, 2, 1, 2,
       1, 1, 2, 3, 1, 3, 1, 2, 2, 3, 1, 2, 2, 2, 1, 3, 1, 3])

预期结果 #

array([ 6. ,  6.5,  8. , 10.5, 12. , 11. , 11.5, 11.5,  6.5, 13.5, 25. ,
       27.5, 30.5, 34.5, 37.5, 36. , 35. , 35. , 34. , 34.5, 34. , 36.5,
       33. , 34. , 34.5, 34.5, 36. , 39. , 37. , 36. , 37. , 36.5, 37.5,
       39. , 36.5, 37.5, 34. , 31. , 27.5, 23. ])

计算逻辑 #

以预期结果第四个元素10.5为例：

1. 取A的第四个元素0.33333333，遍历d中的规则：
   • 规则1（权重1，偏移(-0.25,-0.20)）：计算区间0.33333333-0.25=0.08333333到0.33333333-0.20=0.13333333，A中第二个元素0.11111111落在该区间，对应B值为2，贡献1*2=2。
   • 规则2（权重0.5，偏移(-0.20,-0.10)）：区间0.33333333-0.20=0.13333333到0.33333333-0.10=0.23333333，A中第三个元素0.22222222落在该区间，对应B值为1，贡献0.5*1=0.5。
   • 规则3（权重2，偏移(-0.10,0.15)）：区间0.33333333-0.10=0.23333333到0.33333333+0.15=0.48333333，A中第四、五个元素0.33333333、0.44444444落在该区间，对应B值为2、2，贡献2*(2+2)=8。
2. 累加所有贡献：2+0.5+8=10.5，与预期结果一致。

现有实现（内存效率低） #

D = np.zeros(len(A))
for v in d:
    weight, (_lower, _upper) = v
    lower, upper = A + _lower, A + _upper
    _A = np.tile(A, (len(A), 1))
    __A = np.bitwise_and(_A > lower.reshape(-1, 1), _A < upper.reshape(-1, 1))
    D += weight * (__A @ B)
D

该实现生成了n×n的二维数组_A和__A，当n较大时内存占用会急剧上升。

优化实现方案 #

利用A严格递增的特性，结合np.searchsorted快速定位区间索引，再通过前缀和计算区间内B的和，避免生成大尺寸二维数组：

import numpy as np

# 预处理B的前缀和，方便快速计算区间和
prefix_B = np.concatenate([[0], np.cumsum(B)])
D = np.zeros(len(A))

for weight, (delta_low, delta_high) in d:
    # 计算每个A元素对应的区间左右边界值
    left_vals = A + delta_low
    right_vals = A + delta_high
    
    # 用searchsorted找到每个边界在A中的插入位置，确定区间的左右索引
    left_idx = np.searchsorted(A, left_vals, side='left')
    right_idx = np.searchsorted(A, right_vals, side='right')
    
    # 计算每个区间内B的元素和，前缀和相减即可
    interval_sum = prefix_B[right_idx] - prefix_B[left_idx]
    
    # 加权累加到结果中
    D += weight * interval_sum

# 验证结果（与预期误差可忽略）
print(np.allclose(D, np.array([ 6. ,  6.5,  8. , 10.5, 12. , 11. , 11.5, 11.5,  6.5, 13.5, 25. ,
       27.5, 30.5, 34.5, 37.5, 36. , 35. , 35. , 34. , 34.5, 34. , 36.5,
       33. , 34. , 34.5, 34.5, 36. , 39. , 37. , 36. , 37. , 36.5, 37.5,
       39. , 36.5, 37.5, 34. , 31. , 27.5, 23. ])))  # 输出True

优化优势 #

• 内存优化：仅使用O(n)的内存存储前缀和和结果数组，避免了O(n²)的二维数组开销。
• 性能提升：np.searchsorted是O(n log n)的操作，前缀和计算是O(n)，整体时间复杂度远低于原实现的O(n²)。
• 代码简洁：逻辑更清晰，避免了复杂的矩阵操作。

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