np.random.choice 具有很大的概率数组

Question

我知道我们可以使用概率数组来选择 function，但我的问题是它如何适用于大 arrays。假设我想要 0-65535 之间的 1000 个随机数。 我们如何定义概率数组，使其对于小于 1000 的数字具有 p=0.4，对于 rest 具有 p=0.6？

我试图将数字范围传递给选项 function，但显然，它不是那样工作的。

Answer 1

从文档中，参数p的每个元素都给出了a中对应元素的概率。

由于p和a需要具有相同的大小，因此创建一个与 a 大小相同a p ：

a = np.arange(65536)
n_elem = len(a)

p = np.zeros_like(a, dtype=float)

现在，找到a小于1000的所有元素，并将这些索引的p设置为 0.4 除以小于 1000 的元素数。对于这种情况，您可以对该计算进行硬编码，因为您知道arange的哪些元素更少大于 1000：

p[:1000] = 0.4 / 1000
p[1000:] = 0.6 / 64536

对于a不是从arange派生的一般情况，您可以这样做：

lt1k = a < 1000
n_lt1k = lt1k.sum()

p[lt1k] = 0.4 / n_lt1k
p[~lt1k] = 0.6 / (n_elem - n_lt1k)

请注意， p的总和必须为1 ：

assert np.allclose(p.sum(), 1.0)

现在在choice中使用a和p ：

selection = np.random.choice(a, size=(1000,), p=p)

为了验证选择值 < 1000 的概率是 40%，我们可以检查有多少小于 1000：

print((selection < 1000).sum() / len(selection)) # should print a number close to 0.4

Answer 2

另一种方法是将其视为两种分布的混合：一种以概率 = 0.4 均匀地从 {0..999} 抽取，另一种以概率 = 0.6 均匀地从 {1000..65535} 抽取。

对混合组件使用choice是有道理的，但随后我会使用其他东西来绘制值，因为当概率传递给choice时，它会在每次调用时 O( len(p) ) 工作以转换它们。 Generator.integers应该更有效，因为它可以直接对统一值进行采样。

把这些放在一起，我建议使用类似的东西：

import numpy as np

rng = np.random.default_rng()

n = 1000
splits = np.array([0, 1000, 65536])

# draw weighted mixture components
s = rng.choice(2, n, p=[0.4, 0.6])
# draw uniform values according to component
result = rng.integers(splits[s], splits[s+1])

您可以通过评估np.mean(result < 1000)并检查它是否“接近”0.4 来验证这是从正确的分布中提取的。 其方差约为0.4*0.6 / n ，因此，对于n=1000 ， [0.37, 0.43] 中的值应该在 95% 的时间内可见。

当max(splits) - min(splits)变大时，此方法应该保持快速，而 Pranav 直接使用choice的解决方案会变慢。