Search code examples
pythonnumpysumvectorization

Use numpy to sum indices based on another numpy vector

example problem

I am trying to sum specific indices per row in a numpy matrix, based on values in a second numpy vector. For example, in the image, there is the matrix A and the vector of indices inds. Here I want to sum:

A[0, inds[0]] + A[1, inds[1]] + A[2, inds[2]] + A[3, inds[3]]

I am currently using a python for loop, making the code quite slow. Is there a way to do this using vectorisation? Thanks!

Solution

  • Yes, numpy's magic indexing can do this. Just generate a range for the 1st dimension and use your coords for the second:

    import numpy as np

x1 = np.array( [[1,2,3,4],[5,6,7,8],[9,10,11,12],[13,14,15,16]] )
print(x1[ [0,1,2,3],[2,0,3,1] ].sum())