Numpy在索引对之间子数组的值求和

Question

假设我有一个数组A。然后我有一系列的索引对，比如(a1, b1), (a2, b2) ... (an, bn)

我想要计算这些索引对之间所有元素的和。也就是说，

sum(A[a1:b1]), sum(A[a2:b2]), sum(A[a3:b3]) ...

从运行时间的角度来看，最有效的方法是什么呢？

谢谢！

Answer 1

首先，我会尝试直接的解决办法：

[np.sum(A[a:b]) for (a,b) in ab]

这里的 ab 是一系列的配对。

A[a:b] 是在数组上创建一个视图；这并不涉及数据的复制。

如果这样做还是太慢，请告诉我们更多关于 A 的信息，比如它的大小、你预计会得到多少对索引、还有 (a,b) 的范围是否经常重叠等等。

Answer 2

这里还有另一种方法：

a = np.random.rand(3000)
indices = np.array([[0,3], [9,20], [5,30], [9,33]])
sums = np.add.reduceat(a, indices.ravel())[::2]

assert np.all(sums == np.array([a[i:j].sum() for i,j in indices]))

上面提到的 cumsum 方法，如果有很多索引的话，可能会更高效。

Answer 3

假设你的索引对存储在一个形状为 (n, 2) 的 NumPy 数组 indices 中，并且 n 的值比较大，那么最好避免使用任何 Python 循环。

c = numpy.r_[0, A.cumsum()][indices]
sums = c[:,1] - c[:,0]