Index 2D numpy array by a 2D array of indices without loops

Question

I am looking for a vectorized way to index a numpy.array by numpy.array of indices.

For example:

import numpy as np

a = n         


        

Answer 1

When using arrays of indices to index another array, the shape of each index array should match the shape of the output array. You want the column indices to match inds, and you want the row indices to match the row of the output, something like:

array([[0, 0],
       [1, 1],
       [2, 2]])

You can just use a single column of the above, due to broadcasting, so you can use np.arange(3)[:,None] is the vertical arange because None inserts a new axis:

>>> np.arange(3)[:, None]
array([[0],
       [1],
       [2]])

Finally, together:

>>> a[np.arange(3)[:,None], inds]
array([[0, 3],   # a[0,[0,1]]
       [6, 0],   # a[1,[1,2]] 
       [0, 9]])  # a[2,[0,2]]