非固定样本量下使用Numpy Random Choice生成负二项分布和的优化

方案1：利用负二项分布求和性质（性能最优，推荐） #

负二项分布具备可加性：k个独立同分布的NB(r, p)随机变量之和服从NB(k*r, p)分布，不需要预生成样本池也不需要循环，直接向量化生成结果，性能远高于循环实现，且完全保留default_rng的随机生成特性：

import numpy as np
from numpy.random import default_rng
rng = default_rng()

# 可批量修改测试不同的n、p参数
n_param = 1
p_param = 0.5
# 直接转为numpy数组，支持1e6长度的输入
gmc = np.array([12, 35, 4, 67, 2]) 

# 一行生成结果，无Python层循环
n_pp = rng.negative_binomial(gmc * n_param, p_param)

该方案对1e6长度的gmc数组也能在毫秒级完成运算，完全满足多组参数测试的性能需求。

方案2：需要从固定预生成的nbd样本池抽样的场景 #

如果业务逻辑要求必须从预先生成的固定nbd样本池中做有放回抽样求和，可以用向量化的np.add.reduceat分组求和消除循环：

import numpy as np
from numpy.random import default_rng
rng = default_rng()

nbd = rng.negative_binomial(1, 0.5, int(1e6))
gmc = np.array([12, 35, 4, 67, 2])
total_sample = gmc.sum()

# 一次性抽取所有需要的样本
all_choices = rng.choice(nbd, total_sample)
# 按gmc的长度分组求和
split_indices = np.cumsum(gmc)[:-1]
n_pp = np.add.reduceat(all_choices, np.r_[0, split_indices])

该方案同样完全消除了Python层循环，性能比原生循环提升两个数量级以上。

内容的提问来源于stack exchange，提问作者HEP N008