How to make numpy.argmax return all occurrences of the maximum?
Solution 1:
As documentation of np.argmax
says: "In case of multiple occurrences of the maximum values, the indices corresponding to the first occurrence are returned.", so you will need another strategy.
One option you have is using np.argwhere
in combination with np.amax
:
>>> import numpy as np
>>> listy = [7, 6, 5, 7, 6, 7, 6, 6, 6, 4, 5, 6]
>>> winner = np.argwhere(listy == np.amax(listy))
>>> print(winner)
[[0]
[3]
[5]]
>>> print(winner.flatten().tolist()) # if you want it as a list
[0, 3, 5]
Solution 2:
In case it matters, the following algorithm runs in O(n) instead of O(2n) (i.e., using np.argmax
and then np.argwhere
):
def allmax(a):
if len(a) == 0:
return []
all_ = [0]
max_ = a[0]
for i in range(1, len(a)):
if a[i] > max_:
all_ = [i]
max_ = a[i]
elif a[i] == max_:
all_.append(i)
return all_
Solution 3:
It is even easier, when compared to other answers, if you use np.flatnonzero
:
>>> import numpy as np
>>> your_list = np.asarray([7, 6, 5, 7, 6, 7, 6, 6, 6, 4, 5, 6])
>>> winners = np.flatnonzero(your_list == np.max(your_list))
>>> winners
array([0, 3, 5])
If you want a list:
>>> winners.tolist()
[0, 3, 5]