How to make numpy.argmax return all occurrences of the maximum? How to make numpy.argmax return all occurrences of the maximum? numpy numpy

How to make numpy.argmax return all occurrences of the maximum?


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]


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_


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))>>> winnersarray([0, 3, 5])

If you want a list:

>>> winners.tolist()[0, 3, 5]