A weighted version of random.choice

Мне нужно было написать взвешенную версию random.choice (каждый элемент в списке имеет разную вероятность быть выбранным). Вот что я придумал.:

def weightedChoice(choices):
    """Like random.choice, but each element can have a different chance of
    being selected.

    choices can be any iterable containing iterables with two items each.
    Technically, they can have more than two items, the rest will just be
    ignored.  The first item is the thing being chosen, the second item is
    its weight.  The weights can be any numeric values, what matters is the
    relative differences between them.
    """
    space = {}
    current = 0
    for choice, weight in choices:
        if weight > 0:
            space[current] = choice
            current += weight
    rand = random.uniform(0, current)
    for key in sorted(space.keys() + [current]):
        if rand < key:
            return choice
        choice = space[key]
    return None

Эта функция кажется мне чрезмерно сложной и уродливой. Я надеюсь, что все присутствующие могут предложить несколько предложений по ее улучшению или альтернативные способы сделать это. Эффективность для меня не так важна, как чистота и удобочитаемость кода.

Начиная с версии 1.7.0, в NumPy есть choice функция, поддерживающая распределения вероятностей.

from numpy.random import choice
draw = choice(list_of_candidates, number_of_items_to_pick,
              p=probability_distribution)

Обратите внимание, что probability_distribution это последовательность в том же порядке, что и list_of_candidates. Вы также можете использовать ключевое слово replace=False, чтобы изменить поведение, чтобы нарисованные элементы не заменялись.

Начиная с Python 3.6, существует метод choices из random модуля.

In [1]: import random

In [2]: random.choices(
...:     population=[['a','b'], ['b','a'], ['c','b']],
...:     weights=[0.2, 0.2, 0.6],
...:     k=10
...: )

Out[2]:
[['c', 'b'],
 ['c', 'b'],
 ['b', 'a'],
 ['c', 'b'],
 ['c', 'b'],
 ['b', 'a'],
 ['c', 'b'],
 ['b', 'a'],
 ['c', 'b'],
 ['c', 'b']]

Обратите внимание, что random.choices будет выполнена выборка с заменой в соответствии с документами:

Возвращает k список элементов, выбранных из совокупности, с заменой.

Обратите внимание на полноту ответа:

Когда единица выборки берется из конечной совокупности и возвращается в эту совокупность после того, как были записаны ее характеристики, перед тем, как будет выбрана следующая единица, выборка называется "с заменой". По сути, это означает, что каждый элемент может быть выбран более одного раза.

Если вам нужно выполнить выборку без замены, то, как говорится в блестящем ответе@ronan-paixão, вы можете использовать numpy.choice, аргумент которого replace управляет таким поведением.

def weighted_choice(choices):
   total = sum(w for c, w in choices)
   r = random.uniform(0, total)
   upto = 0
   for c, w in choices:
      if upto + w >= r:
         return c
      upto += w
   assert False, "Shouldn't get here"

1. Распределите веса в кумулятивное распределение.


2. Используйте random.random() для выбора случайного значения с плавающей точкой 0.0 <= x < total.


3. Выполните поиск в дистрибутиве с помощью bisect.разделите пополам, как показано в примере на http://docs.python.org/dev/library/bisect.html#other-examples.

from random import random
from bisect import bisect

def weighted_choice(choices):
    values, weights = zip(*choices)
    total = 0
    cum_weights = []
    for w in weights:
        total += w
        cum_weights.append(total)
    x = random() * total
    i = bisect(cum_weights, x)
    return values[i]

>>> weighted_choice([("WHITE",90), ("RED",8), ("GREEN",2)])
'WHITE'

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