在numpy中随机选择元素

Question

我有一个元素概率的数组，比如说 [0.1, 0.2, 0.5, 0.2]。这个数组的总和是1.0。

我想用普通的Python或者numpy来根据这些概率来随机选择元素：第一个元素大约10%的时间被选中，第二个20%，第三个50%等等。这个“选择”应该返回被选中元素的索引。

我想出了这个方法：

def draw(probs):
    cumsum = numpy.cumsum(probs / sum(probs)) # sum up to 1.0, just in case
    return len(numpy.where(numpy.random.rand() >= cumsum)[0])

这个方法可以用，但感觉太复杂了，应该有更简单的办法。谢谢。

Answer 1

使用 numpy.random.multinomial - 这是最有效的方法

Answer 2

你想从一个分类分布中抽样，但在numpy这个库里没有直接实现。不过，多项分布是分类分布的一个扩展，可以用来实现这个功能。

>>> import numpy as np
>>> 
>>> def sampleCategory(p):
...     return np.flatnonzero( np.random.multinomial(1,p,1) )[0]
... 
>>> sampleCategory( [0.1,0.5,0.4] )
1

Answer 3

---

import numpy as np
def random_pick(choices, probs):
    '''
    >>> a = ['Hit', 'Out']
    >>> b = [.3, .7]
    >>> random_pick(a,b)
    '''
    cutoffs = np.cumsum(probs)
    idx = cutoffs.searchsorted(np.random.uniform(0, cutoffs[-1]))
    return choices[idx]

In [22]: import numpy as np
In [23]: probs = [0.1, 0.2, 0.5, 0.2]

In [24]: cutoffs = np.cumsum(probs)
In [25]: cutoffs
Out[25]: array([ 0.1,  0.3,  0.8,  1. ])

In [26]: np.random.uniform(0, cutoffs[-1])
Out[26]: 0.9723114393023948

In [27]: cutoffs.searchsorted(0.9723114393023948)
Out[27]: 3

它是怎么工作的：

计算累积和：

在半开区间 [0, cutoffs[-1]) 中生成一个均匀分布的随机数：

使用 searchsorted 来找到这个随机数应该插入到 cutoffs 的位置索引：

返回 choices[idx]，其中 idx 就是那个索引。