python - 用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

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)