Get elements from a numpy array using another array-CodePudding

I have a numpy array, a, of shape (2, 5, 2). I need to get 2 elements from the array at locations [0, 1, 0] and [1, 3, 0]. To retrieve these elements, I have an array, b, with entries [1, 3]. When I use a[:,b], I get 2x2 matrix (as expected). How can I change the indexing to return the required vector instead? Following is an MWE.

import numpy as np

a = np.random.randn(2, 5, 2)
b = np.random.randint(0, 5, 2)

c = a[:, b]
c.shape  # prints (2, 2), while I only need a[0, b[0]] and a[1, b[1]].

PS: This problem can be solved using a for loop. However in my actual code the index is in the order of 1000 and therefore not feasible.

CodePudding user response：

As pointed out in the comments, one can do c = a[np.arange(a.shape[0]), b] to obtain the required result.