Numpy: Fancy Indexing over Multiple arrays

Question

Is there an efficient way to index over multiple arrays?

For example, I have an array I want to index from

a = [[1,2,3],[4,5,6]]

And another array contains the indices. b = [[0, 1], [1,2]]

And I expect [[1, 2], [5, 6]], which indexes the first row of a by [0,1], and indexes the second row of a by [1,2].

Thanks.

Answer 1

In [107]: a = [[1,2,3],[4,5,6]]
In [108]: b = [[0, 1], [1,2]]

a and b are lists. The appropriate solution is a nested list comprehension

In [111]: [[a[i][j] for j in x] for i,x in enumerate(b)]
Out[111]: [[1, 2], [5, 6]]

Now if a is made into a numpy array:

In [112]: np.array(a)[np.arange(2)[:,None], b]
Out[112]: 
array([[1, 2],
       [5, 6]])

For this the 1st dimension of the array is indexed with a (2,1) array, and the 2nd with a (2,2). They broadcast together to produce a (2,2) result.

Numpy extract submatrix

is working in the same direction, but the accepted answer use ix_

Y[np.ix_([0,3],[0,3])]

which won't work in the case of a (2,2) b.

In [113]: np.array(a)[np.ix_(np.arange(2), b)]
ValueError: Cross index must be 1 dimensional

ix_ will turn the 1st dimension np.arange(2) in to the right (2,1).

---

This might make the broadcasting more explicit:

In [114]: np.array(a)[[[0,0],[1,1]], [[0,1],[1,2]]]
Out[114]: 
array([[1, 2],
       [5, 6]])

It selects elements (0,0), (0,1), (1,1) and (1,2)

---

To further test this, make b non symmetic:

In [138]: b = [[0, 1,1], [1,2,0]]       # (2,3)
In [139]: np.array(a)[np.arange(2)[:,None], b]
Out[139]: 
array([[1, 2, 2],
       [5, 6, 4]])