python - 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!

Answer 1

You can use advanced indexing:

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

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

Or with np.take_along_axis:

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

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