如何无需逐行遍历计算Numpy数组每行指定索引区间的均值？

问题描述 #

给定两个形状为m的索引数组idx0和idx1，以及形状为m×n的二维数组d，需要计算每行中从idx0[i]到idx1[i]（左闭右开）区间内元素的均值，要求避免逐行遍历，寻求最快实现方式。

示例代码 #

import numpy as np

idx0 = np.array([1, 3, 2, 5])
idx1 = np.array([5, 8, 6, 7])
d = np.array([[1,2,3,4,5,6,7,8,9], 
              [1,2,3,4,5,6,7,8,9], 
              [1,2,3,4,5,6,7,8,9], 
              [1,2,3,4,5,6,7,8,9]])

期望结果 #

np.array([3.5, 6, 4.5, 6.5])

最优解决方案：前缀和矢量化计算 #

无需逐行遍历，利用前缀和数组结合numpy的矢量化索引即可高效实现，这是大规模数据场景下速度最快的方案之一。

实现步骤 #

1. 计算数组d的行前缀和，为兼容idx0=0的情况，在每行开头补0；
2. 利用数组索引提取每行idx1和idx0位置的前缀和，求差得到区间元素总和；
3. 用总和除以区间长度（idx1 - idx0）得到均值。

代码实现 #

import numpy as np

# 计算行前缀和，每行开头补0以处理idx0=0的情况
prefix = np.pad(np.cumsum(d, axis=1), ((0, 0), (1, 0)), mode='constant')

# 生成行索引，提取对应区间的前缀和差值
row_indices = np.arange(d.shape[0])
interval_sums = prefix[row_indices, idx1] - prefix[row_indices, idx0]

# 计算最终均值
means = interval_sums / (idx1 - idx0)

print(means)
# 输出：array([3.5, 6. , 4.5, 6.5])

方法对比 #

• 列表推导：实现简单但本质是Python级循环，当m达到10^5级别的大规模数据时，速度会远慢于矢量化操作；
• Numba加速循环：能提升循环效率，但存在编译启动开销，仍不如纯numpy矢量化操作的执行效率；
• 前缀和矢量化：完全基于numpy底层C实现，无Python循环开销，是大规模数据场景下的最优选择。

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