python - 从python中的权重数组中获取随机索引的快速方法

Question

我经常发现自己需要一个数组或列表的随机索引，其中索引的概率不是均匀分布的，而是根据某些正权重。获得它们的快速方法是什么？我知道我可以将权重作为可选参数 p 传递给 numpy.random.choice，但是这个函数看起来很慢，并且构建一个 arange 到通过它也不理想。权重之和可以是任意正数并且不保证为 1，这使得在 (0,1] 中生成随机数然后减去权重条目直到结果为 0 或更小的方法是不可能的。

虽然有关于如何以简单的方式实现类似的东西(主要不是关于获取数组索引，而是对应的元素)的答案，例如Weighted choice short and simple ，我正在寻找一个快速的解决方案，因为经常执行适当的功能。我的权重经常变化，所以构建类似别名掩码的开销(详细介绍可以在 http://www.keithschwarz.com/darts-dice-coins/ 上找到)应该被视为计算时间的一部分。

Answer 1

累加和平分

在任何一般情况下，计算权重的累积总和似乎是可取的，并使用 bisect 模块中的 bisect 在生成的排序数组中找到一个随机点

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

如果速度是一个问题。下面给出更详细的分析。

注意:如果数组不是平面的，numpy.unravel_index 可用于将平面索引转换为形状索引，如https://stackoverflow.com/a/19760118/1274613 所示。

实验分析

使用 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 https://stackoverflow.com/a/19760118/1274613
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 方法也比 type-naive bisect变体。 (这些结果是使用 Python 3.3.5 和 numpy 1.8.1 获得的，因此其他版本可能有所不同。)基于 numpy.random.multinomial 的函数对于大权重的效率低于基于累积求和的方法。据推测，argmax 必须遍历整个数组并在每一步进行比较这一事实起着重要作用，从增加和减少权重列表之间的四秒差异也可以看出这一点。