Slice of 2d numpy array with another array-CodePudding

I have a quite large 2d array, and I need to get both the index of the maximum value in axis 1, and the maximum value itself. I can retrieve these two values as follows:

import numpy as np
a = np.arange(27).reshape(9, 3)
idx = np.argmax(a, axis=1)
max_val = np.max(a, axis=1)

However, since I have already found the index of the maximum value, it feels like I should be able to construct the array of maximum values using idx without having to look up the value again.

I realise I can use np.choose(idx, a.T) but this involves transposing the matrix which will be much more expensive than just using max. I can do something like np.array([a[i][idx[i]] for i in range(len(a))]) but this involves creating a list which again seems more expensive that just calling np.max.

Is there any way to slice a with idx in numpy without restructuring the array?

CodePudding user response：

Your a and argmax:

In [602]: a
Out[602]: 
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]])

In [603]: idx
Out[603]: array([2, 2, 2, 2, 2, 2, 2, 2, 2], dtype=int64)

A common way of using that index array:

In [606]: a[np.arange(a.shape[0]),idx]
Out[606]: array([ 2,  5,  8, 11, 14, 17, 20, 23, 26])

A newer tool, that may be easier to use (if not familiar with the first):

In [607]: np.take_along_axis(a,idx[:,None],1)
Out[607]: 
array([[ 2],
       [ 5],
       [ 8],
       [11],
       [14],
       [17],
       [20],
       [23],
       [26]])