Question

我有两个数组，值和索引

>>> values
array([[5, 4, 2, 4, 6],
       [7, 9, 7, 3, 6]])
>>> indexes
array([[2, 4],
       [0, 3],
       [0, 1],
       [1, 3]])

我想要的是一种快速的方法（因为我的数组非常大），为每个值获取与索引中所有索引集合对应的元素总和。

我希望，第一个值[5,4,2,4,6]得到

>>> values[0][indexes.flatten()].reshape(indexes.shape)
array([[2, 6],
       [5, 4],
       [5, 4],
       [4, 4]])

>>> values[0][indexes.flatten()].reshape(indexes.shape).sum(axis=1)
array([8, 9, 9, 8])

使用这种技术并循环遍历所有值是我能想到的最快的。有没有更好的办法？提前感谢您的时间。

Answer 1

方法＃1

只需索引到列并沿最后一个轴求和 -

values[:,indexes].sum(axis=-1)

示例运行 -

In [39]: values
Out[39]: 
array([[5, 4, 2, 4, 6],
       [7, 9, 7, 3, 6]])

In [40]: indexes
Out[40]: 
array([[2, 4],
       [0, 3],
       [0, 1],
       [1, 3]])

In [41]: values[:,indexes].sum(axis=-1)
Out[41]: 
array([[ 8,  9,  9,  8],
       [13, 10, 16, 12]])

方法＃2

如果indexes的每一行都没有重复项，我们只需使用matrix-multiplication即可获得sum-reductions，这会更快 -

m,n = indexes.shape[0], values.shape[1]
mask = np.zeros((n,m),dtype=bool) # faster with float dtype
mask[indexes, np.arange(m)[:,None]] = 1
out = values.dot(mask)

从带有索引的数组中选择值并对它们求和

1 个答案: