Numpy矢量化指数的总和

时间:2016-10-11 14:41:41

标签: python numpy vectorization

我有一个索引列表(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]

1 个答案:

答案 0 :(得分:2)

这是使用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(),但计算上应该可以忽略不计。