非规则数组大小的Numpy随机选择

Question

我正在从负二项分布(nbd(中生成一个随机选择的和的数组，每个和的长度都不规则。现在我实现如下：

import numpy
from numpy.random import default_rng
rng = default_rng()
nbd = rng.negative_binomial(1, 0.5, int(1e6))
gmc = [12, 35, 4, 67, 2]
n_pp = np.empty(len(gmc))
for i in range(len(gmc)):
n_pp[i] = np.sum(rng.choice(nbd, gmc[i]))

这是可行的，但当我在实际数据上执行它时，速度非常慢(gmc的维度为1e6(，我想对nbd中的n和p的多个值进行更改(在本例中，它们分别设置为1和0.5(。

我想找出一种蟒蛇般的方法来做到这一点，消除循环，但我不确定这是否可能。如果可能的话，我想保留default_rng，以获得比旧的生成方式更好的随机生成(np.srandom.cachoice(。

Answer 1

参数为(n, p)的负二项分布的m样本之和的分布是参数为(m*n, p)的负二项式分布。因此，您可以直接使用negative_binomial(gmc, 0.5):生成结果，而不是对来自预先计算的negative_binomial(1, 0.5)大样本的随机选择进行求和

In [68]: gmc = [12, 35, 4, 67, 2]
In [69]: npp = rng.negative_binomial(gmc, 0.5)
In [70]: npp
Out[70]: array([ 9, 34,  1, 72,  7])

(negative_binomial方法将广播其输入，因此我们可以将gmc作为参数传递，以通过一次调用生成所有样本。(

更一般地说，如果要更改用于生成nbd的n，则需要将该n乘以gmc中的相应元素，然后将乘积传递给rng.negative_binomial。