How to get indexes of maximum elements in NumPy array of arrays?

Question

Python

There is an array of the following form:
(the largest elements whose indices must be obtained are marked in red)

That is, for such an array, the code should return something like this:
(0, 2, 2), (1, 0, 1), (2, 1 , 0), (3, 0, 0), (4, 0, 0)
Of course, this can be done by looping with np.argmax, but perhaps someone can find a more efficient solution implemented only with numpy (in my In this case, the speed of the script is very important)
Perhaps these functions will help: np.argmax, np.unravel_index

UPD:

I found a way to get the elements themselves, but here's the problem, I need to get their indices

Answer 1

my array

import numpy as np
x = np.random.random((5,3,3)) * 10 // 1

array([,

       ,

       ,

       ,

       ])

Since the maximum is sought among the 2nd level, we can change the dimension, considering each nine as a series:

y = x.reshape(5,9).argmax(axis=1)

# array([0, 7, 5, 2, 2])

Now these indices need to be converted back to a pair of coordinates:

f = lambda i: (i // 3, i % 3)
# или вернее так:
f = lambda i: (i // x.shape[1], i % x.shape[2])

(a, b) = f(y)

# array([0, 2, 1, 0, 0]), array([0, 1, 2, 2, 2])

It remains to pull out pairs with the same indices from both arrays.

result = np.empty((a.size + b.size,), dtype=a.dtype)
result[0::2] = a
result[1::2] = b
result.reshape(5,2)

array()