Numpy如何以向量化方式基于另一数组索引计算整数数组切片和

核心思路 #

利用Numpy前缀和（累加和）的特性实现完全向量化计算，无需任何Python层面的循环，计算效率极高，适配任意大小的输入数组。
核心逻辑：sum(a[x:y]) = 前缀和[y] - 前缀和[x]，只需构造前缀和数组，再对对应索引位做差即可得到结果。

完整实现代码 #

import numpy as np

# 示例输入
a = np.array([10,10,10,8,8,8])
b = np.array([0,2,3,5])

# 构造前缀和数组，首位补0，prefix_sum[i] = sum(a[:i])
prefix_sum = np.concatenate([[0], np.cumsum(a)])
# 拼接0到b数组头部，再求相邻位的差值，即为结果
result = np.diff(prefix_sum[np.insert(b, 0, 0)])

print(result)
# 输出：[ 0 20 10 16]

代码说明 #

• 前缀和数组首位补0是为了直接对应sum(a[:0])=0的计算逻辑，省去单独处理边界的步骤
• np.insert(b, 0, 0)将0插入b数组头部，构造出所有区间的起止索引序列
• np.diff直接对前缀和的对应索引值做差，一步得到所有区间的求和结果
• 如果偏好更简洁的语法，可以用np.r_替代拼接操作，写法如下：

  prefix_sum = np.r_[0, a.cumsum()]
  result = np.diff(prefix_sum[np.r_[0, b]])
  

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