具有自定义概率分布的随机索引

Question

我有一个概率向量（当然总和为 1）：

prob = [0.1, 0.3, 0.4, 0.2]

现在我需要为该向量生成一个随机索引（包括 0 到 3 之间的数字），但我希望每个索引的概率由 prob

给出

0 will be generated with prob 0.1
1 will be generated with prob 0.3
2 will be generated with prob 0.4
3 will be generated with prob 0.2

我知道我可以通过计算 cumsum

cumsum = [0.1, 0.4, 0.8, 1.0]

然后生成一个0到1之间的随机数：

rand_num = np.random.random()

最后使用 np.digitize 检查我的随机数落在哪个 bin 中。

idx = np.digitize([rand_num], cumsum)

这行得通，我对此很满意，digitize 甚至接受数字列表并将它们分类到 bin 中，因此我可以创建自己的函数来生成给定概率分布的索引。

我的问题是：这是一个常见问题，所以是否已经存在执行此操作的函数？（这比自己做更有效率）

谢谢

Answer 1

您可以使用 random.choices 从 Python 3.6 开始，其中包括一个 weights 参数：

>>> from random import choices
>>> prob = [0.1, 0.3, 0.4, 0.2]
>>> choices(range(len(prob)), weights=prob)
[2]
>>> choices(range(len(prob)), weights=prob)
[3]
>>> choices(range(len(prob)), weights=prob, k=4)
[1, 2, 2, 2]

具有自定义概率分布的随机索引

Random index with custom probability distribution

python

random

numpy

distribution