numpy using multidimensional index array on another multidimensional array

Question

I have a 2 multidimensional arrays, and I'd like to use one as the index to produce a new multidimensional array. For example:

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

b = array([[0,2],[1,1],[3,1]])

I want to use the first array in b to return those indexed elements in the first array of a, and so on. So I want the output to be:

array([[4,2],[8,8],[6,1]])

This is probably simple but I couldn't find an answer by searching. Thanks.

Answer 1

This is a little tricky, but the following will do it:

>>> a[np.arange(3)[:, np.newaxis], b]
array([[4, 2],
       [8, 8],
       [6, 1]])

You need to index both the rows and the columns of the a array, so to match your b array you would need an array like this:

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

And then a[rows, b] would clearly return what you are after. You can get the same result relying on broadcasting as above, replacing the rows array with np.arange(3)[:, np.newaxis], which is equivalent to np.arange(3).reshape(3, 1).