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

【问题标题】：Numpy: Computing sum(array[i, a[i]:b[i]]) for all rows iNumpy：计算所有行 i 的 sum(array[i, a[i]:b[i]])
【发布时间】：2023-03-09 21:48:01
【问题描述】：

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

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

仅使用 numpy 是否有更智能/更快的方法来做到这一点？

【问题讨论】：

我想你会想要扁平化 arr 并使用 numpy.add.reduceat 来计算总和。
这有一个不幸的效果，即用arr 的其他部分计算一堆你不需要的总和。我不知道在不编写显式循环的情况下避免这种情况的好方法。

标签： python arrays numpy sum slice

【解决方案1】：

这是使用NumPy broadcasting 和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)

【讨论】：

我喜欢这个答案！