Aditya369 - 1 month ago 35
Python Question

# Numpy vectorize sum over indices

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]
``````

Answer

Here's an almost* vectorized approach using `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()]
``````

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.

Source (Stackoverflow)
Comments