在numpy中绘制随机元素

Question

我想有一系列元素概率[0.1, 0.2, 0.5, 0.2].该数组总和为1.0.

使用普通的Python或numpy,我想绘制与其概率成比例的元素:第一个元素约10%的时间,第二个20%,第三个50%等."draw"应该返回绘制元素的索引.

我想出了这个:

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

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

在半开区间内计算均匀分布的随机数[0, cutoffs[-1]):

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

使用searchsorted查找将插入随机数的索引cutoffs:

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

返回choices[idx],idx该索引在哪里.