如何用数组操作而非循环对Numpy数组第0列做当前元素与后三行元素逐元素求和？

针对百万级行数的数组，必须用矢量化操作替代循环来保证效率，以下是两种可行的方案：

方案1：使用滑动窗口视图（sliding_window_view） #

这个方法直观易懂，直接生成滑动窗口后求和，适合需要明确窗口结构的场景：

import numpy as np

# 生成示例数组（真实场景替换为你的百万行数组）
k = np.random.rand(10, 1) * 10
k = k.astype(int)

# 创建滑动窗口：每个窗口包含当前行+后续3行，窗口形状为(4,1)
windows = np.lib.stride_tricks.sliding_window_view(k, window_shape=(4, 1))
# 对每个窗口内的所有元素求和
sum_result = windows.sum(axis=(1, 2))

# 生成与原数组行数一致的结果，不足窗口长度的位置补NaN（可根据需求替换为0或其他值）
full_result = np.full_like(k, np.nan)
full_result[:len(sum_result)] = sum_result.reshape(-1, 1)

print("原数组：")
print(k)
print("\n求和结果（当前元素+后3行，末尾不足补NaN）：")
print(full_result)

方案2：使用卷积运算（np.convolve） #

卷积运算底层优化程度极高，对于一维求和场景速度更快：

import numpy as np

k = np.random.rand(10, 1) * 10
k = k.astype(int).ravel()  # 转成一维数组简化卷积操作

# 用全1卷积核实现窗口求和，mode='valid'只保留完全匹配的窗口结果
sum_result = np.convolve(k, np.ones(4, dtype=int), mode='valid')

# 还原为二维数组并补全长度
full_result = np.full_like(k, np.nan).reshape(-1, 1)
full_result[:len(sum_result)] = sum_result.reshape(-1, 1)

关键说明 #

• 两种方法均为numpy底层C实现的矢量化操作，处理百万行数组的速度远快于Python循环；
• 若不需要补全结果长度，直接取sum_result即可，它包含所有能取到完整4个元素的窗口求和值；
• 若后续行数不足3行（比如最后3行），可根据业务需求选择补NaN、0，或者截断结果。

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