从具有索引的数组中选择值并对其求和

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)