Use numpy to sum indices based on another numpy vector-CodePudding

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!

CodePudding user response：

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())