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] |
这是一个使用
的几乎*矢量化方法
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()] |
请注意,这有一些设置开销,因此它最适合具有大量列的
*:几乎是因为在开始时使用了