Extract values from one array using the indices of maximum values of another array-CodePudding

I have two 2D arrays x and y.

x = np.array([[2,4,6],
              [9,4,6],
              [6,8,3]])

y = np.array([[88,55,33],
              [43,87,65],
              [98,34,56]])

Using the argmax function, I found the indices of the maximum values of x along axis 1.

idx = x.argmax(axis=1)

output: array([2, 0, 1], dtype=int64)

now, I want the the values from array y, which are on these specific indices. My expected array is ([33,43,34]).

I tried using y[idx] but it gives the following output.

array([[98, 34, 56],
   [88, 55, 33],
   [43, 87, 65]])

How can I get the output I want?

CodePudding user response：

Easy as using np.arange :

result = y[np.arange(len(idx)), idx]