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()，但在计算上应该可以忽略不计。