Use numpy argmax on the 3rd dimension of numpy array

Question

I have a 3d numpy array of shape (150, 9, 5) and I would like to get the argmax of the 3rd dimension (out of the elements) but I cannot get it done. For example, let's say the array looks like (this example does not follow the shape I am using in my code):

[[[615, 142, 18]],

   [[12, 412, 98]],

   [[5, 11, 162]],

   [[19, 76, 7]],

   [[56, 78, 12]],

   [[119, 41, 55]],

   [[26, 7, 15]],

   [[19, 67, 53]]]

I would like to get back an array of shape (150, 9, 1) (again this is in my case, not related to the example). For the example it would be:

[[[0]],

   [[1]],

   [[2]],

   [[1]],

   [[1]],

   [[0]],

   [[0]],

   [[1]]]

When I try np.argmax() with both axis 0 and 1, I get wrong results.

Is there a way to work it out directly or should I go through each (9,5) using a for loop ?

`ind[...,None]` with `ind` being the argmax indices? – Divakar May 03 '19 at 16:35 — Divakar, May 03 '19 at 16:35

score 0 · Accepted Answer · answered May 03 '19 at 16:37

0

You've got extra axes with extra dimensions:

import numpy as np
a=np.array([[[615, 142, 18]],

   [[12, 412, 98]],

   [[5, 11, 162]],

   [[19, 76, 7]],

   [[56, 78, 12]],

   [[119, 41, 55]],

   [[26, 7, 15]],

   [[19, 67, 53]]])

a.argmax(axis=2)

Output:

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

answered May 03 '19 at 16:37

Juan C

5,846
2
17
51

Thanks a lot !! This does the job – Wazaki May 03 '19 at 16:41
You're welcome! Could you please mark it as the accepted answer so there's no confusion? Thanks ! – Juan C May 03 '19 at 16:47
Sure I will ! I need to wait some time before I can do that, which is why I haven't yet (SO rules) – Wazaki May 03 '19 at 16:48
Thanks @Wazaki ! – Juan C May 03 '19 at 19:21

Use numpy argmax on the 3rd dimension of numpy array

1 Answers1