numpy的：计算总和（阵列[I，A [1]：B [I]]）对于所有i行

我有一个numpy的阵列arr和切片的列表开始点start和切片端点end的列表。对于每行i，我想确定从start[i]到end[i]的元素总和。也就是说，我想确定numpy的：计算总和（阵列[I，A [1]：B [I]]）对于所有i行

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

是否有更聪明/更快的方式来使用numpy来做到这一点？

来源

2016-09-28 Samufi

我想你会想'扁平化和arr'使用['numpy.add.reduceat'（HTTP：//文档。 scipy.org/doc/numpy/reference/generated/numpy.ufunc.reduceat.html）来计算总和。 – user2357112

不过，计算一堆你不需要的'arr'的其他部分会带来不幸的效果。我不知道避免这种情况的好方法，不用写一个明确的循环。 – user2357112

下面是使用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)

来源

2016-09-28 20:53:33 Divakar

我喜欢这个答案！ – Kasramvd

numpy的：计算总和（阵列[I，A [1]：B [I]]）对于所有i行

回答

相关问题