Tricky numpy argmax on last dimension of 3-dimensional ndarray

Question

if have an array of shape (9,1,3).

array([[[  6,  12, 108]],

   [[122, 112,  38]],

   [[ 57, 101,  62]],

   [[119,  76, 177]],

   [[ 46,  62,   2]],

   [[127,  61, 155]],

   [[  5,   6, 151]],

   [[  5,   8, 185]],

   [[109, 167,  33]]])

I want to find the argmax index of the third dimension, in this case it would be 185, so index 7.

I guess the solution is linked to reshaping but I can't wrap my head around it. Thanks for any help!

Answer 1

You may have to do it like this:

data = np.array([[[  6,  12, 108]],

   [[122, 112,  38]],

   [[ 57, 101,  62]],

   [[119,  76, 177]],

   [[ 46,  62,   2]],

   [[127,  61, 155]],

   [[  5,   6, 151]],

   [[  5,   8, 185]],

   [[109, 167,  33]]])

np.argmax(data[:,0][:,2])
 7

Answer 2

I'm not sure what's tricky about it. But, one way to get the index of the greatest element along the last axis would be by using np.max and np.argmax like:

# find `max` element along last axis 
# and get the index using `argmax` where `arr` is your array
In [53]: np.argmax(np.max(arr, axis=2))
Out[53]: 7

Alternatively, as @PaulPanzer suggested in his comments , you could use:

In [63]: np.unravel_index(np.argmax(arr), arr.shape)
Out[63]: (7, 0, 2)

In [64]: arr[(7, 0, 2)]
Out[64]: 185