Numpy：计算 sum(array[i, a[i]:b[i]]) 所有行 i

Question

我有一个 numpy 数组 arr 和一个切片起点列表 start 以及一个切片端点列表 end。对于每一行 i，我想确定从 start[i] 到 end[i] 的元素之和。也就是我要确定

[np.sum(arr[i, start[i]:end[i]]) for i in range(arr.shape[0])]

是否有 smarter/faster 仅使用 numpy 的方法？

Answer 1

这是使用 NumPy broadcasting and np.einsum -

的矢量化方法

# Create range array corresponding to the length of the no. of cols
r = np.arange(arr.shape[1])

# Mask of ranges corresponding to the start and end indices using broadcasting
mask = (start[:,None] <= r) & (end[:,None] > r)

# Finally, we use the mask to select and sum rows using einsum
out = np.einsum('ij,ij->i',arr,mask)

Numpy：计算 sum(array[i, a[i]:b[i]]) 所有行 i

Numpy: Computing sum(array[i, a[i]:b[i]]) for all rows i

python

arrays

numpy

sum

slice