Быстрый способ получить случайный индекс из массива весов в python

Я регулярно нахожусь в положении необходимости случайного индекса для массива или списка, где вероятности индексов неравномерно распределены, но в соответствии с определенными положительными весами. Какой быстрый способ получить их? Я знаю, что могу передать весы numpy.random.choice в качестве необязательного аргумента p, но функция кажется довольно медленной, и создание arange для прохождения не идеально. Сумма весов может быть произвольным положительным числом и не гарантируется равной 1, что делает подход к генерации случайного числа в (0,1), а затем вычитанию весовых позиций до тех пор, пока результат не окажется 0 или менее невозможным.

Хотя есть ответы на то, как реализовать подобные вещи (в основном не о получении индекса массива, а о соответствующем элементе) простым способом, например Взвешенный выбор короткого и простого Я ищу быстрое решение, потому что соответствующая функция выполняется очень часто. Мои веса меняются часто, поэтому накладные расходы на создание чего-то вроде маски псевдонима (подробное введение можно найти на http://www.keithschwarz.com/darts-dice-coins/) следует учитывать часть времени вычисления.

Ответы

Ответ 1

Кумулятивное суммирование и деление пополам

В любом общем случае представляется целесообразным вычислить суммарную сумму весов и использовать bisect из модуля bisect, чтобы найти случайную точку в результирующем отсортированном массиве

def weighted_choice(weights):
    cs = numpy.cumsum(weights)
    return bisect.bisect(cs, numpy.random.random() * cs[-1])

если скорость вызывает беспокойство. Более подробный анализ приведен ниже.

Примечание. Если массив не является плоским, numpy.unravel_index можно использовать для преобразования плоского индекса в форменный индекс, как показано в fooobar.com/info/28987/...

Экспериментальный анализ

Существует четыре более или менее очевидных решения с использованием встроенных функций numpy. Сравнение всех из них с помощью timeit дает следующий результат:

import timeit

weighted_choice_functions = [
"""import numpy
wc = lambda weights: numpy.random.choice(
    range(len(weights)),
    p=weights/weights.sum())
""",
"""import numpy
# Adapted from /info/28987/efficient-way-of-sampling-from-indices-of-a-numpy-array/212123#212123
def wc(weights):
    cs = numpy.cumsum(weights)
    return cs.searchsorted(numpy.random.random() * cs[-1], 'right')
""",
"""import numpy, bisect
# Using bisect mentioned in https://stackoverflow.com/a/13052108/1274613
def wc(weights):
    cs = numpy.cumsum(weights)
    return bisect.bisect(cs, numpy.random.random() * cs[-1])
""",
"""import numpy
wc = lambda weights: numpy.random.multinomial(
    1,
    weights/weights.sum()).argmax()
"""]

for setup in weighted_choice_functions:
    for ps in ["numpy.ones(40)",
               "numpy.arange(10)",
               "numpy.arange(200)",
               "numpy.arange(199,-1,-1)",
               "numpy.arange(4000)"]:
        timeit.timeit("wc(%s)"%ps, setup=setup)
    print()

Результирующий результат

178.45797914802097
161.72161589498864
223.53492237901082
224.80936180002755
1901.6298267539823

15.197789980040397
19.985687876993325
20.795070077001583
20.919113760988694
41.6509403079981

14.240949985047337
17.335801470966544
19.433710905024782
19.52205040602712
35.60536142199999

26.6195822560112
20.501282756973524
31.271995796996634
27.20013752405066
243.09768892999273

Это означает, что numpy.random.choice на удивление очень медленный, и даже выделенный метод numpy searchsorted медленнее, чем вариант наивного типа bisect. (Эти результаты были получены с использованием Python 3.3.5 с numpy 1.8.1, поэтому для других версий могут быть разные.) Функция, основанная на numpy.random.multinomial, менее эффективна для больших весов, чем методы, основанные на суммарном суммировании. Предположительно, тот факт, что argmax должен перебирать весь массив и выполнять сравнения, каждый шаг играет значительную роль, что также видно из четырех секундного разницы между увеличивающимся и уменьшающимся списком веса.