Numpy argmax - random breakdown

Question

In numpy.argmax , the break function between multiple maximum elements is such that the first element is returned. Is there functionality for randomizing a communication break so that all maximum numbers have an equal probability of choice?

The following is an example directly from the 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 for the 1st and 5th items in the list to return with equal probability.

Thanks!

Answer 1

Use np.random.choice -

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

Let the array with three maximum candidates be checked -

 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 will not work.

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 the generation of random floats is slower than any other method, random.random can be replaced with another method.

Answer 3

The easiest way is

 np.random.choice(np.where(b == b.max())[0])