How Can I Index This Numpy Array?

Given a numpy matrix a of shape (5,3), and a index vector b of shape (5,), where each entry in the index vector is between 0 to 2, how can I create a new vector c based on a and it

Solution 1:

Use arange for the other dimension:

c = a[np.arange(5), b]