I have a list of indices (list(int)) and a list of summing indices (list(list(int)). Given a 2D numpy array, I need to find the sum over indices in the second list for each column and add them to the corresponding indices in the first column. Is there any way to vectorize this?
Here is the normal code:
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]
Here's an almost* vectorized approach using
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()]
Please note that this has some setup overhead, so it would be best suited for
2D input arrays with a good number of columns as we are iterating along the columns.
*: Almost because of the use of
map() at the start, but computationally that should be negligible.