如何用Numpy实现基于B值筛选的A数组部分滚动求和（替代循环）

问题背景 #

现有两个一维Numpy数组A和B，代表一段时期内按天索引的价格数据，生成方式示例如下：

import random
import numpy as np

A = np.array([random.randint(1,100) for _ in range(1000)])
B = np.array([random.randint(1,100) for _ in range(1000)])

需求：对A执行窗口大小为20的部分滚动求和——针对每一天，取过去20天中B值较大的10天，求和对应位置的A值。原Python循环实现效率极低，需改用Numpy原生高效实现。

原循环实现（存在笔误，sum(window[1][10:])应为sum(x[1] for x in window[10:])）：

result = np.zeros(A.shape)
for i in range(19, len(A)):
    window = list(zip(B[i - 19: i + 1], A[i - 19: i + 1]))
    window.sort(key = lambda x: x[0])
    result[i] = sum(x[1] for x in window[10:])

Numpy高效实现方案 #

利用Numpy的向量化运算和滑动窗口工具，完全避免Python循环，大幅提升运行效率：

import numpy as np
import random

# 1. 生成测试数据
A = np.array([random.randint(1,100) for _ in range(1000)])
B = np.array([random.randint(1,100) for _ in range(1000)])

# 2. 生成A和B的滑动窗口（窗口大小20）
# 生成的窗口数组shape为(981, 20)，对应从第20天开始的每个窗口
A_windows = np.lib.stride_tricks.sliding_window_view(A, window_shape=20)
B_windows = np.lib.stride_tricks.sliding_window_view(B, window_shape=20)

# 3. 对每个窗口内的B值排序，取B值最大的10个元素的索引
# argsort返回从小到大的索引，取后10个即为B值最大的10个位置
sorted_indices = np.argsort(B_windows, axis=1)
top10_indices = sorted_indices[:, -10:]

# 4. 取出对应A值并求和
top10_A_sum = np.take_along_axis(A_windows, top10_indices, axis=1).sum(axis=1)

# 5. 构建最终结果数组（前19天无足够窗口数据，设为0）
result = np.zeros_like(A)
result[19:] = top10_A_sum

效率说明 #

• sliding_window_view基于内存 stride 机制实现，无需复制原始数据，内存占用极低
• 所有核心操作均为Numpy底层优化的向量化运算，相比Python循环，在1000长度数组上可实现数十倍甚至百倍的速度提升，数组规模越大优势越明显

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