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!

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]) 

Answer 2

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.