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!

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

Solution 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.

Solution 3

Easiest way is

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

Solution 4

Since the accepted answer may not be obvious, here is how it works:

• b == b.max() will return an array of booleans, with values of true where items are max and values of false for other items
• flatnonzero() will do two things: ignore the false values (nonzero part) then return indices of true values. In other words, you get an array with indices of items matching the max value
• Finally, you pick a random index from the array

Solution 5

Additional to @Manux's answer,

Changing b.max() to np.amax(b,**kw, keepdims=True) will let you do it along axes.

def randargmax(b,**kw):
    """ a random tie-breaking argmax"""
    return np.argmax(np.random.random(b.shape) * (b==b.max()), **kw)

randargmax(b,axis=None) 

View more solutions

Comments

[crypdick]

if you have floats instead of ints, you may want to replace b == b.max() with np.isclose(b, b.max())

[Jaymin]

Please provide some explanation along with the answer.

[ijuneja]

arr.max() can be replaced by np.max(arr) for clarity that numpy array is involved

[Mehdi Abbassi]

roughly three times slower!

[Mehdi Abbassi]

roughly 18 times slower!