python - Numpy 向量化索引之和

Question

我有一个索引列表 (list(int)) 和一个求和索引列表 (list(list(int))。给定一个 2D numpy 数组，我需要在第二个列表中找到每个索引的总和列并将它们添加到第一列中的相应索引中。有什么方法可以对其进行矢量化吗？这是正常的代码:

indices = [1,0,2]
summing_indices = [[5,6,7],[6,7,8],[4,5]]
matrix = np.arange(9*3).reshape((9,3))
for c,i in enumerate(indices):
    matrix[i,c] = matrix[summing_indices[i],c].sum()+matrix[i,c]

Answer 1

这是使用 np.add.reduceat 的几乎*矢量化方法-

lens = np.array(map(len,summing_indices))
col = np.repeat(indices,lens)
row = np.concatenate(summing_indices)
vals = matrix[row,col]
addvals = np.add.reduceat(vals,np.append(0,lens.cumsum()[:-1]))
matrix[indices,np.arange(len(indices))] += addvals[indices.argsort()]

请注意，这有一些设置开销，因此它最适合具有大量列的2D输入数组，因为我们沿着列进行迭代。

*:几乎是因为在开始时使用了 map()，但从计算角度来看，这应该可以忽略不计。