Numpy argmax - случайная разбивка

В numpy.argmax функция разрыва между несколькими максимальными элементами такова, что возвращается первый элемент. Существует ли функциональность для рандомизации разрыва связи так, что все максимальные числа имеют равную вероятность выбора?

Ниже приведен пример непосредственно из документации numpy.argmax.

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

Я ищу способы, чтобы 1-й и 5-й элементы в списке возвращались с равной вероятностью.

Спасибо!

Ответ 1

Используйте np.random.choice -

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

Пусть проверяется массив с тремя максимальными кандидатами -

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

Ответ 2

В случае многомерного массива choice не будет работать.

Альтернативой является

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

Если по какой-либо причине генерация случайных поплавков медленнее, чем какой-либо другой метод, random.random может быть заменен другим методом.

Ответ 3

Самый простой способ это

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