How to index numpy array?

Question

A simple example: I have a numpy array([[5, 4, 2],[1, 8, 8],[4, 5, 6]]) and another array([2, 2, 1]). The second array contains indices of interest for each row in the first array. I want my final array to be simply a vector array([2,8,5]). I'm looking for something along the lines of first_array[second_array] but it is proving elusive. Thank you!

Answer 1

You could try something like the following:

>>> arr
array([[5, 4, 2],
       [1, 8, 8],
       [4, 5, 6]])
>>> idx
array([2, 2, 1])
>>> arr[np.arange(arr.shape[0]), idx]
array([2, 8, 5])