Python生成截断负二项分布内核崩溃，如何优化代码提升运行效率？

崩溃原因 #

你原代码采用的拒绝采样逻辑在截断阈值远小于分布主体区间时会出现极低的采样效率。你当前测试参数下，scipy负二项分布（参数意义为成功n次前的失败次数）的均值为 n*(1-p)/p = 231，远高于你设置的max_value=99，符合条件的样本占比不足3%，拒绝采样过程中temp_size会指数级增长，很快生成超量样本占满内存，导致内核崩溃。

优化方案 #

推荐使用逆变换采样替代拒绝采样，完全避免无效采样的内存开销，优化后代码如下：

import scipy.stats as sct
import numpy as np

def truncated_Nbinom(n, p, max_value, size):
    # 计算截断范围内的所有整数值对应的CDF
    x_range = np.arange(0, max_value + 1)
    cdf_vals = sct.nbinom.cdf(x_range, n, p)
    # 归一化CDF到[0,1]区间，对应截断分布的CDF
    normalized_cdf = cdf_vals / cdf_vals[-1]
    # 生成均匀随机数，映射为对应分布的样本
    random_u = np.random.uniform(0, 1, size=size)
    samples = np.searchsorted(normalized_cdf, random_u)
    return samples

效果验证 #

你给出的测试调用语句可以直接正常运行，无内核崩溃问题：

input_1 = truncated_Nbinom(99, 0.3, 99, 5000).tolist()

该方案时间复杂度稳定为O(max_value + size)，不受截断阈值位置影响，不会出现内存暴涨问题。

内容的提问来源于stack exchange，提问作者Gözde Filiz