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

arrays, numpy, python

Solution

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]]

Problem

I am looking for a vectorized way to index a `numpy.array` by `numpy.array` of indices. For example: ``` import numpy as np a = np.array([[0,3,4], [5,6,0], [0,1,9]]) inds = np.array([[0,1], [1,2], [0,2]]) ``` I want to build a new array, such that every row(i) in that array is a row(i) of array `a`, indexed by row of array inds(i). My desired output is: ``` array([[ 0., 3.], # a[0][:,[0,1]] [ 6., 0.], # a[1][:,[1,2]] [ 0., 9.]]) # a[2][:,[0,2]] ``` I can achieve this with a loop: ``` def loop_way(my_array, my_indices): new_array = np.empty(my_indices.shape) for i in xrange(len(my_indices)): new_array[i, :] = my_array[i][:, my_indices[i]] return new_array ``` But I am looking for a pure vectorized solution.

Original source