获得n个独立几何随机变量之和的最快方法

Question

我有很大的整数n和几何分布：

n = 1000000
distribution = scipy.stats.geom(0.53284, loc=-1)

我想计算n此分布的独立随机变量的总和。我可以这样做：

distribution.rvs(size=n).sum()

但这需要很长时间。什么是更快的方法呢？

Answer 1

独立几何随机变量的总和具有负二项式分布。

因此，请使用scipy.stats.nbinom：

import numpy as np
import scipy.stats as stats
import matplotlib.pyplot as plt

n = 10**4
p = 0.53284
nsamples = 10**3
distribution = stats.geom(p, loc=-1)
x = distribution.rvs(size=(nsamples, n))
total = x.sum(axis=1)

distribution2 = stats.nbinom(n, p)
total2 = distribution2.rvs(nsamples)

fig, ax = plt.subplots(nrows=2, sharex=True)
ax[0].hist(total, normed=True)
ax[0].set_title('by summing geom')
ax[1].hist(total2, normed=True)
ax[1].set_title('by using nbinom')
plt.show()

对于上面的示例，使用nbinom速度提高了约448倍：

In [64]: %timeit distribution = stats.geom(p, loc=-1); distribution.rvs(size=(nsamples, n)).sum(axis=1)
1 loop, best of 3: 397 ms per loop

In [65]: %timeit distribution2 = stats.nbinom(n, p); total2 = distribution2.rvs(nsamples)
1000 loops, best of 3: 885 µs per loop

In [66]: 397000/885
Out[66]: 448.5875706214689