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!
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])
In the case of a multi-dimensional array, choice won't work.
An alternative is
def randargmax(b,**kw):
""" a random tie-breaking argmax"""
return np.argmax(np.random.random(b.shape) * (b==b.max()), **kw)
If for some reason generating random floats is slower than some other method, random.random can be replaced with that other method.
Easiest way is
np.random.choice(np.where(b == b.max())[0])