在python中生成截断的负二项分布

import numpy as np
from scipy.stats import nbinom

def truncated_negative_binomial(n, p, lower, upper, size):
    # 生成截断的负二项分布
    dist = nbinom(n, p)
    samples = []
    while len(samples) < size:
        sample = dist.rvs()
        if lower <= sample <= upper:
            samples.append(sample)
    return samples

# 设置参数
n = 5  # 成功次数
p = 0.3  # 成功概率
lower = 2  # 截断下界
upper = 10  # 截断上界
size = 1000  # 生成样本数量

# 生成截断的负二项分布样本
samples = truncated_negative_binomial(n, p, lower, upper, size)

# 打印样本均值和方差
mean = np.mean(samples)
variance = np.var(samples)
print("样本均值：", mean)
print("样本方差：", variance)