关于python:Numpy对索引进行矢量化总和

Numpy vectorize sum over indices

我有一个索引列表 (list(int)) 和一个求和索引列表 (list(list(int))。给定一个 2D numpy 数组,我需要在第二个列表中找到每个索引的总和列并将它们添加到第一列中的相应索引.有没有办法向量化这个?
这是正常的代码:

1
2
3
4
5
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]

这是一个使用 np.add.reduceat -

的几乎*矢量化方法

1
2
3
4
5
6
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(),但在计算上应该可以忽略不计。