使用Python计算二项分布概率矩阵

Question

给定 N 和 P，我希望得到一个二维二项分布概率矩阵 M。

for i in range(1, N+1):
   for j in range(i+1):
      M[i,j] = choose(i, j) * p**j * (1-p)**(i-j)
other value = 0

我想知道是否有快速的方法来获取这个矩阵，而不是使用for循环。变量N可能大于100,000。

- Albert Feng

---

看一下这个链接：https://docs.scipy.org/doc/numpy-1.14.0/reference/generated/numpy.random.binomial.html 你可能不需要整个矩阵... - ma3oun

@ma3oun 谢谢，但我需要整个矩阵。 - Albert Feng

如果您需要生成整个矩阵，最低的时间复杂度是O(n^2)，这就是您for循环的性能。没有任何解决方案可以比您现有的速度提高数量级。 - Him

@Scott 确实如此，但时间复杂度是一个理论上与语言无关的概念。底层实现也很重要。 - Brad Solomon

@BradSolomon 我并不是在暗示问题没有答案，只是在设定期望。 :) - Him

即具体实现取决于大量的因素。一些关于理论上限的简单计算也很重要。 - Him

Answer 1

我相信 scipy.stats.binom 可以按照你所需要的方式利用广播。

# Binomial PMF: Pr(X=k) = choose(n, k) * p**k * (1-p)**(n-k)
# Probability of getting exactly k successes in n trials

>>> from scipy.stats import binom

>>> n = np.arange(1, N+1, dtype=np.int64)
>>> dist = binom(p=0.25, n=n)
>>> M = dist.pmf(k=np.arange(N+1, dtype=np.int64)[:, None])

>>> M.round(2)
array([[0.75, 0.56, 0.42, 0.32, 0.24, 0.18, 0.13, 0.1 , 0.08, 0.06],
       [0.25, 0.38, 0.42, 0.42, 0.4 , 0.36, 0.31, 0.27, 0.23, 0.19],
       [0.  , 0.06, 0.14, 0.21, 0.26, 0.3 , 0.31, 0.31, 0.3 , 0.28],
       [0.  , 0.  , 0.02, 0.05, 0.09, 0.13, 0.17, 0.21, 0.23, 0.25],
       [0.  , 0.  , 0.  , 0.  , 0.01, 0.03, 0.06, 0.09, 0.12, 0.15],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.02, 0.04, 0.06],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.01, 0.02],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ],
       [0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  , 0.  ]])

这里，行数是k（从0开始计数），列数是n（从1开始计数）：

>>> from math import factorial as fac
>>> def manual_pmf(p, n, k):
...     return fac(n) / (fac(k) * fac(n - k)) * p**k * (1-p)**(n-k)

>>> manual_pmf(p=0.25, n=3, k=2)
0.140625  # (2, 2) in M because M's columns are effectively 1-indexed 

您也可以从零开始设置n，以获得在行和列上都是0索引的数组：

>>> n = np.arange(N+1, dtype=np.int64)  # M.shape is (11, 11)

- Brad Solomon