efficient per column matrix indexing in numpy

Question

I have two matrices of the same size, A, B. I want to use the columns of B to acsses the columns of A, on a per column basis. For example,

A = np.array([[1, 4, 7],
              [2, 5, 8],
              [3, 6, 9]])

and

B = np.array([[0, 0, 2],
              [1, 2, 1],
              [2, 1, 0]])

I want something like:

A[B] = [[1, 4, 9],  
        [2, 6, 8], 
        [3, 5, 7]]

I.e., I've used the j'th column of B as indices to the j'th column of A. Is there any effiecnt way of doing so? Thanks!

score 3 · Accepted Answer · answered Apr 12 '20 at 13:41

3

A[B, np.arange(A.shape[0])]

array([[1, 4, 9],
       [2, 6, 8],
       [3, 5, 7]])

np.take_along_axis(A, B, axis=0)

array([[1, 4, 9],
       [2, 6, 8],
       [3, 5, 7]])

answered Apr 12 '20 at 13:41

yatu

Amazing works great. Just to make sure I understand what is happening in the first suggestion: you index A at axis 0 with the index Bij, while at axis 1 the vector np.arange(A.shape[0]) is broadcasted to the same shape as B? – MRm Apr 12 '20 at 13:55
1

Exactly @MRm Internally both indexing arrays are broadcast together to the resulting shape. The arange is just specifying which columns each value in the rows of `B` correspond to. Find a detailed explanation along with examples [here](https://towardsdatascience.com/numpy-indexing-explained-c376abb2440d) – yatu Apr 12 '20 at 14:00

1 Answers1