numpy: how to get a max from an argmax result

Question

I have a numpy array of arbitrary shape, e.g.:

a = array([[[ 1,  2],
            [ 3,  4],
            [ 8,  6]],

          [[ 7,  8],
           [ 9,  8],
         


        

Answer 1

For arbitrary-shape arrays, the following should work :)

a = np.arange(5 * 4 * 3).reshape((5,4,3))

# for last axis
argmax = a.argmax(axis=-1)
a[tuple(np.indices(a.shape[:-1])) + (argmax,)]

# for other axis (eg. axis=1)
argmax = a.argmax(axis=1)
idx = list(np.indices(a.shape[:1]+a.shape[2:]))
idx[1:1] = [argmax]
a[tuple(idx)]

or

a = np.arange(5 * 4 * 3).reshape((5,4,3))

argmax = a.argmax(axis=0)
np.choose(argmax, np.moveaxis(a, 0, 0))

argmax = a.argmax(axis=1)
np.choose(argmax, np.moveaxis(a, 1, 0))

argmax = a.argmax(axis=2)
np.choose(argmax, np.moveaxis(a, 2, 0))

argmax = a.argmax(axis=-1)
np.choose(argmax, np.moveaxis(a, -1, 0))