为何纯Python实现的负二项分布pmf比Scipy版本快得多？

scipy.stats提供的nbinom.pmf()函数用于计算负二项分布的概率质量函数（PMF），其数学逻辑可以用一段极简的纯Python代码实现：

from math import comb
def nbinom_pmf(k, n, p):
        return comb(k+n-1, n-1)* p**n * (1-p)**k

实际测试显示，Scipy的实现比上述纯Python代码慢不少，即使使用官方推荐的._pmf跳过部分外层检查，多数场景下仍慢于纯Python版本：

In [24]: %timeit nbinom_pmf(1, 26, 0.5)
282 ns ± 1.61 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

In [25]: %timeit nbinom._pmf(1, 26, 0.5)
2.03 µs ± 6.55 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [32]: %timeit nbinom._pmf(36, 26, 0.5)
2.03 µs ± 1.49 ns per loop (mean ± std. dev. of 7 runs, 100,000 loops each)

In [33]: %timeit nbinom_pmf(36, 26, 0.5)
1.64 µs ± 30.9 ns per loop (mean ± std. dev. of 7 runs, 1,000,000 loops each)

导致这种性能差异的核心原因在于两者的设计目标和实现复杂度不同：

• 通用场景适配开销：Scipy的统计函数是为通用场景设计的，需要支持多维数组输入、numpy广播机制、多种数据类型兼容，还要适配整个stats模块的统一接口，这些都会带来额外的逻辑分支和调用开销。而纯Python函数只处理标量输入，逻辑完全极简。
• 参数校验与边界处理：Scipy的PMF实现会对输入的k、n、p做严格的合法性校验——比如检查k是否为非负整数、n是否为正、p是否在0到1区间内，还要处理p=0或p=1等边界情况。即使使用._pmf跳过了部分外层检查，内部仍保留了必要的校验和类型转换步骤，而纯Python函数默认输入都是合法的，完全跳过了这些流程。
• 数值稳定性优化：Scipy的实现会考虑极端参数下的数值稳定性，比如当参数很大时，直接计算组合数和幂次容易出现溢出或精度损失，Scipy会通过对数转换等方法规避这类问题，这些额外的数值处理会增加计算时间。纯Python实现直接计算，在小参数场景下更快，但在极端参数下可能出现精度问题。
• 多层调用链损耗：Scipy的nbinom.pmf()涉及多层函数调用，从公共API到内部私有方法，还要对接numpy的底层逻辑，即使._pmf也不是最底层的实现。而纯Python函数直接调用math.comb，调用链极短，额外损耗几乎可以忽略。

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