CODEWITHSUNDEEP
pythonnumpyvectorization
Aditya369
Aditya369

Reputation: 834

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]

Upvotes: 3

Views: 1747

Answers (1)

Divakar
Divakar

Reputation: 221564

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.

Upvotes: 3

Related Questions