numpy ndarray advanced indexing

Question

I have ndarray of 3 dimension. How do I select index 0 and 1 from first axis while selecting index 0 and 3 from second axis and index 1 from third axis?

I tried to use index [(0,1), (1, 3), 1], which produces a result completely different than what I thought it would produce.

So two questions here. What does [(0,1), (1, 3), 1] do? And how to correctly create an index that solve my original question.

a = np.arange(30).reshape(3, 5, 2)
array([[[ 0,  1],
        [ 2,  3],
        [ 4,  5],
        [ 6,  7],
        [ 8,  9]],

       [[10, 11],
        [12, 13],
        [14, 15],
        [16, 17],
        [18, 19]],

       [[20, 21],
        [22, 23],
        [24, 25],
        [26, 27],
        [28, 29]]])

a[0, (1, 3), 1]  # produces array([3, 7])
a[(0,1), (1, 3), 1] # produces array([ 3, 17])

```

Answer 1

When you index the way you're doing it, NumPy doesn't interpret it as selecting those indices of each dimension. Instead, NumPy broadcasts the arguments against each other:

a[(0,1), (1, 3), 1] -> a[array([0, 1]), array([1, 3]), array([1, 1])]

and then creates a result array where a[i, j, k][x] == a[i[x], j[x], k[x]] .

To get the behavior you're looking for, you need to reshape the arguments you're passing, so that broadcasting them against each other produces an array of shape (2, 2) instead of shape (2,) . This means the first argument needs shape (2, 1) , the second argument needs to have shape (1, 2) or (2,) , and the third argument's shape is fine. numpy.ix_ can make this easier, but it doesn't support scalar arguments. a[np.ix_([0, 1], [1, 3], [1])] does what you would have expected a[[0, 1], [1, 3], [1]] to do, but to get the shape you would have expected from a[[0, 1], [1, 3], 1] , your options are messier:

>>> a[np.ix_([0, 1], [1, 3], [1])]
array([[[ 3],
        [ 7]],

       [[13],
        [17]]])
>>> a[np.ix_([0, 1], [1, 3]) + (1,)]
array([[ 3,  7],
       [13, 17]])
>>> a[np.ix_([0, 1], [1, 3], [1])][:, :, 0]
array([[ 3,  7],
       [13, 17]])