Use numpy to sum indices based on another numpy vector

Question

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!

Tim Roberts · Accepted Answer · 2022-02-13 20:00:31Z

1

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

answered Feb 13, 2022 at 20:00

Tim Roberts

55.3k4 gold badges29 silver badges41 bronze badges

Sign up to request clarification or add additional context in comments.

Collectives™ on Stack Overflow

Use numpy to sum indices based on another numpy vector

1 Answer 1

Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

Comments

Your Answer

Sign up or log in

Post as a guest

Related