在Python中生成符合给定PDF的随机数

Question

我想生成一个整数的随机数，并且这个随机数的出现概率是根据一个列表来分配的。比如说，如果我有一个列表pdf=[3,2,1]，那么我希望调用 rndWDist(pdf) 能返回0、1和2，出现的概率分别是3/6、2/6和1/6。因为在随机模块里找不到这样的功能，所以我自己写了一个函数。

def randintWDist(pdf):
    cdf=[]
    for x in pdf:
        if cdf:
            cdf.append(cdf[-1]+x)
        else:
            cdf.append(x)
    a=random.randint(1,cdf[-1])
    i=0
    while cdf[i]<a:
        i=i+1
    return i

有没有更简单的方法可以达到同样的效果呢？

Answer 1

如果你使用的是numpy（版本1.7或更新的版本），你可以使用这个叫做 np.random.choice 的功能：

In [27]: import numpy as np

In [28]: distribution = (1./6, 2./6, 3./6)

In [29]: np.random.choice(np.arange(len(distribution)), p=distribution)
Out[29]: 0

In [30]: np.random.choice(np.arange(len(distribution)), p=distribution, size=10)
Out[30]: array([2, 1, 1, 2, 2, 0, 1, 0, 1, 0])

Answer 2

根据你输入的格式，你可以这样做：

def randint_with_dist(pdf):
    choices = []
    for index, value in enumerate(pdf):
        choices.extend(index for _ in range(value))
    return random.choice(choices)

因为每次传入相同的pdf时都会使用同样的列表，你可以考虑把这个列表缓存起来，这样可以提高效率（不过会占用一些空间）：

def randint_with_dist(pdf, choices={}):
    pdf = tuple(pdf)
    if pdf not in choices:
        choices[pdf] = []
        for index, value in enumerate(pdf):
            choices[pdf].extend(index for _ in range(value))
    return random.choice(choices[pdf])

Answer 3

这是一个重复的问题：生成具有特定（数值）分布的随机数

根据第一个回答的建议，你可以使用scipy.stats.rv_discrete。

你可以这样使用它：

from scipy.stats import rv_discrete
numbers = (1,2,3)
distribution = (1./6, 2./6, 3./6)
random_variable = rv_discrete(values=(numbers,distribution))
random_variable.rvs(size=10)

这会返回一个包含10个随机值的numpy数组。