如何修改现有多项式概率代码以获取正确的三项式展开系数？

已实现的多项式概率函数 #

基于二项式概率函数，扩展实现了计算多项式概率的Python函数：

公式说明 #

• 原二项式概率函数：

  f(n, k, p) = n! / (k! * (n - k)!) * p^k * (1 - p)^(n - k)

• 扩展后的多项式概率函数：

  f(n, k, p) = n! / (k[0]! * k[1]! * ... * k[x]!) * p[0]^k[0] * p[1]^k[1] * ... p[x]^k[x]

实现代码 #

from math import prod


def factorial(n: int) -> int:
    r = 1
    for i in range(1, n + 1):
        r *= i
    return r


def combination(n: int, k: list) -> float:
    return factorial(n) / prod(map(factorial, k))


def multinomial_probability(n: int, k: list, p: list) -> float:
    return combination(n, k) * prod(map(pow, p, k))

二项式展开系数验证 #

该函数可正确获取二项式展开的概率值（系数除以2^n）：

for n in range(5):
    print([multinomial_probability(n=n, k=[k, n - k], p=[1/2, 1/2]) for k in range(n + 1)])

输出结果：

[1.0]
[0.5, 0.5]
[0.25, 0.5, 0.25]
[0.125, 0.375, 0.375, 0.125]
[0.0625, 0.25, 0.375, 0.25, 0.0625]

三项式展开系数获取问题 #

尝试获取三项式展开系数时，多次调用均未得到预期的1、3、5、7、9项结果：

尝试1 #

for n in range(5):
    print([multinomial_probability(n=n, k=[m, k - m, n - k], p=[1/3, 1/3, 1/3]) for k in range(n + 1) for m in range(k + 1)])

生成的列表元素数量为1、3、6、10、15，不符合预期。

尝试2 #

desired_outputs = [1, 3, 5, 7, 9]
for n in range(5):
    print([multinomial_probability(n=n, k=[k, n - 2 * k, k], p=[1/3, 1/3, 1/3]) for k in range(desired_outputs[n])])

尝试3 #

desired_outputs = [1, 3, 5, 7, 9]
for n in range(5):
    print([multinomial_probability(n=n, k=[k, n - k - 1, 1], p=[1/3, 1/3, 1/3]) for k in range(desired_outputs[n])])

上述尝试的输出均不符合预期的三项式展开系数。

解决方案 #

你的multinomial_probability函数本身是正确的，问题出在生成k组合的方式。以下两种场景的实现方式：

场景1：获取标准三项式展开系数（所有项） #

标准三项式展开(a+b+c)^n的项数为(n+1)(n+2)/2，对应所有满足k0 + k1 + k2 = n的非负整数组合，实现代码：

def generate_trinomial_coefficients(n):
    coefficients = []
    for k0 in range(n + 1):
        for k1 in range(n - k0 + 1):
            k2 = n - k0 - k1
            # 传入p=[1,1,1]直接获取系数，无需乘以概率
            coeff = multinomial_probability(n, [k0, k1, k2], [1, 1, 1])
            coefficients.append(coeff)
    return coefficients

for n in range(5):
    print(generate_trinomial_coefficients(n))

输出结果：

[1.0]
[1.0, 1.0, 1.0]
[1.0, 2.0, 1.0, 2.0, 2.0, 1.0]
[1.0, 3.0, 3.0, 1.0, 6.0, 3.0, 3.0, 6.0, 3.0, 1.0]
[1.0, 4.0, 6.0, 4.0, 1.0, 12.0, 12.0, 4.0, 12.0, 12.0, 6.0, 12.0, 6.0, 4.0, 1.0]

场景2：获取对称分组后的三项式系数（1、3、5、7、9项） #

若你需要的是按k0 - k2差值分组后的系数和（项数为2n+1），实现代码：

def generate_symmetric_trinomial_coefficients(n):
    coefficients = []
    # 遍历k0 - k2的差值，从-n到n
    for s in range(-n, n + 1):
        total = 0
        k2 = max(0, -s)
        while True:
            k0 = k2 + s
            k1 = n - k0 - k2
            if k0 < 0 or k1 < 0:
                break
            coeff = multinomial_probability(n, [k0, k1, k2], [1, 1, 1])
            total += coeff
            k2 += 1
        coefficients.append(total)
    return coefficients

for n in range(5):
    print(generate_symmetric_trinomial_coefficients(n))

输出结果：

[1.0]
[1.0, 2.0, 1.0]
[1.0, 4.0, 6.0, 4.0, 1.0]
[1.0, 6.0, 15.0, 20.0, 15.0, 6.0, 1.0]
[1.0, 8.0, 28.0, 56.0, 70.0, 56.0, 28.0, 8.0, 1.0]

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