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

It’s possible, although somewhat non-obvious to do this as follows:

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

The index np.arange(a.shape[0]) simply indexes the rows to which the array of column indices inds applies. Appending [:, None] modifies the shape of this array such that its shape is (a.shape[0], 1), i.e. each row index is in a separate row of a 1-column-wide 2D array.

The basic principle is that the number of dimensions in the index arrays must agree, and their shapes must also do so. See documentation for np.ix_ to get a feel for this.