ChatGPT解决这个技术问题 Extra ChatGPT

Argmax of numpy array returning non-flat indices

I'm trying to get the indices of the maximum element in a Numpy array. This can be done using numpy.argmax. My problem is, that I would like to find the biggest element in the whole array and get the indices of that.

numpy.argmax can be either applied along one axis, which is not what I want, or on the flattened array, which is kind of what I want.

My problem is that using numpy.argmax with axis=None returns the flat index when I want the multi-dimensional index.

I could use divmod to get a non-flat index but this feels ugly. Is there any better way of doing this?


S
Sven Marnach

You could use numpy.unravel_index() on the result of numpy.argmax():

>>> a = numpy.random.random((10, 10))
>>> numpy.unravel_index(a.argmax(), a.shape)
(6, 7)
>>> a[6, 7] == a.max()
True

e
eumiro
np.where(a==a.max())

returns coordinates of the maximum element(s), but has to parse the array twice.

>>> a = np.array(((3,4,5),(0,1,2)))
>>> np.where(a==a.max())
(array([0]), array([2]))

This, comparing to argmax, returns coordinates of all elements equal to the maximum. argmax returns just one of them (np.ones(5).argmax() returns 0).


This will iterate the array three times, not only twice. One time to find the maximum, a second time to build the result of ==, and a third time to extract the True values from this result. Note that there might be more than one item equal to the maximum.
R
Richard Ludlow

To get the non-flat index of all occurrences of the maximum value, you can modify eumiro's answer slightly by using argwhere instead of where:

np.argwhere(a==a.max())

>>> a = np.array([[1,2,4],[4,3,4]])
>>> np.argwhere(a==a.max())
array([[0, 2],
       [1, 0],
       [1, 2]])

It's not effective since you get three passes and a matrix creation. Imagine we've got 9000x7000 image (A3@600dpi) - would you still insist on your solution?