python - Numpy argmax - random tie breaking

Question

In numpy.argmax function, tie breaking between multiple max elements is so that the first element is returned. Is there a functionality for randomizing tie breaking so that all maximum numbers have equal chance of being selected?

Below is an example directly from numpy.argmax documentation.

>>> b = np.arange(6)
>>> b[1] = 5
>>> b
array([0, 5, 2, 3, 4, 5])
>>> np.argmax(b) # Only the first occurrence is returned.
1

I am looking for ways so that 1st and 5th elements in the list are returned with equal probability.

Thank you!

question from:https://stackoverflow.com/questions/42071597/numpy-argmax-random-tie-breaking

Answer 1

Use np.random.choice -

np.random.choice(np.flatnonzero(b == b.max()))

Let's verify for an array with three max candidates -

In [298]: b
Out[298]: array([0, 5, 2, 5, 4, 5])

In [299]: c=[np.random.choice(np.flatnonzero(b == b.max())) for i in range(100000)]

In [300]: np.bincount(c)
Out[300]: array([    0, 33180,     0, 33611,     0, 33209])