How can I calculate the multinomial probability values?

I am trying to write a function that finds the multinomial expansion. In order to do that I modified the binomial probability function;

From:

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

To:

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

The codes are:

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))

To get the coefficients of binomial expansion, I am calling the multinomial_probability function with the following arguments:

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)])

The result I get is below and it's the correct result:

[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]

Now I am thinking how to find the coefficients of the trinomial expansion.

In order to do that, I called the multinomial_probability function with the following arguments, despite I knew that it was not going to give me the correct result.

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)])

The above code printed the following result:

[1.0]
[0.3333333333333333, 0.3333333333333333, 0.3333333333333333]
[0.1111111111111111, 0.2222222222222222, 0.2222222222222222, 0.1111111111111111, 0.2222222222222222, 0.1111111111111111]
[0.03703703703703703, 0.1111111111111111, 0.1111111111111111, 0.1111111111111111, 0.2222222222222222, 0.1111111111111111, 0.03703703703703703, 0.1111111111111111, 0.1111111111111111, 0.03703703703703703]
[0.012345679012345677, 0.0493827160493827, 0.0493827160493827, 0.07407407407407407, 0.14814814814814814, 0.07407407407407407, 0.0493827160493827, 0.14814814814814814, 0.14814814814814814, 0.0493827160493827, 0.012345679012345677, 0.0493827160493827, 0.07407407407407407, 0.0493827160493827, 0.012345679012345677]

The number of elements of the lists I would like to get are 1, 3, 5, 9, 11 in order. The code I wrote produces lists with element numbers of 1, 3, 6, 10, 15, and the probabilities are calculated according to the number of elements of these lists.

So the code I wrote produces lists with element numbers of 1, 3, 6, 10, 15, and the probabilities are calculated according to the number of elements of these lists. We can't give decimal numbers to k or m either. Multiplying two ints does not yield 1, 3, 5, 7, 9 in order.

Now, I am changing the code as below:

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])])

The output I get is like this:

[1.0]
[0.3333333333333333, 0.3333333333333333, 0.08333333333333334]
[0.1111111111111111, 0.2222222222222222, 0.055555555555555566, 0.006172839506172837, 0.00038580246913580245]
[0.03703703703703703, 0.2222222222222222, 0.05555555555555555, 0.0061728395061728366, 0.00038580246913580234, 1.5432098765432092e-05, 4.286694101508915e-07]
[0.012345679012345677, 0.14814814814814814, 0.07407407407407407, 0.008230452674897117, 0.0005144032921810698, 2.0576131687242793e-05, 5.71559213534522e-07, 1.166447374560249e-08, 1.8225740227503888e-10]

This is not correct either.

I change k like this:

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])])

And this produces the following output:

[1.0]
[0.3333333333333333, 0.3333333333333333, 0.16666666666666669]
[0.2222222222222222, 0.2222222222222222, 0.1111111111111111, 0.037037037037037035, 0.009259259259259257]
[0.1111111111111111, 0.2222222222222222, 0.1111111111111111, 0.03703703703703703, 0.009259259259259259, 0.001851851851851851, 0.0003086419753086419]
[0.0493827160493827, 0.14814814814814814, 0.14814814814814814, 0.0493827160493827, 0.012345679012345677, 0.0024691358024691353, 0.0004115226337448558, 5.878894767783656e-05, 7.348618459729569e-06]

What kind of change I should make in order to get the coefficients of the following output (trinomial expansion)?

Thanks in advance...

I think you need a function partitions(n, d) to find all the ways to divide a power n onto d variables.

Then you can calculate multicomb(n, k) for each k in partitions(n, d)

from math import prod, factorial
from itertools import combinations_with_replacement
from string import ascii_lowercase as alphabet
from typing import Iterator

def multicomb(n: int, k: list) -> int:
    '''assume sum(k) == n'''
    return factorial(n) // prod(map(factorial, k))

def diff(k: list) -> tuple:
    '''un-cumsum a list of cumulative sums'''
    result = []
    c = 0
    for x in k:
        result.append(x - c)
        c = x
    return tuple(result)

def partitions(n: int, d: int) -> Iterator[int]:
    for k in combinations_with_replacement(range(n+1), d-1):
        yield diff(k+(n,))

def prettyfactor(c: int) -> str:
    return str(c) if c != 1 else ''

def prettyvar(xi: int, ki: int) -> str:
    if ki == 0:
        return ''
    elif ki == 1:
        return str(xi)
    else:
        return f'{xi}**{ki}'

def product_expansion(n: int, d: int) -> str:
    variables = alphabet[:d] # assume d <= 26
    return ' + '.join(
        prettyfactor(multicomb(n,k)) + ' ' + ' '.join(prettyvar(xi,ki) for xi,ki in zip(variables, k))
        for k in partitions(n, d)
    )

print(product_expansion(2, 2)) # (a + b)**2
  b**2 + 2 a b +  a**2 

print(product_expansion(3, 2)) # (a + b)**3
  b**3 + 3 a b**2 + 3 a**2 b +  a**3 

print(product_expansion(2, 3)) # (a + b + c)**2
   c**2 + 2  b c +   b**2  + 2 a  c + 2 a b  +  a**2  

print(product_expansion(3, 3)) # (a + b + c)**3
   c**3 + 3  b c**2 + 3  b**2 c +   b**3  + 3 a  c**2 + 6 a b c + 3 a b**2  + 3 a**2  c + 3 a**2 b  +  a**3  

I wrote function product_expansion to return a string representing the expanded product (a + b + c + ...)**n.

Instead we can go back to your probabilities question. We want to return the probability of each possible combination. An appropriate return value is a dict, mapping combinations to probabilities.

def weighted_dice_probabilities(n: int, p: list) -> dict:
    '''roll n identical weighted dice
    face i has probability p[i]
    return a dict mapping each possible outcome to its proba'''
    d = len(p)
    return {
        k: multicomb(n, k) * prod(map(pow, p, k))
        for k in partitions(n, d)
    }

print(weighted_dice_probabilities(0, [1/3, 1/3, 1/3]))
{(0, 0, 0): 1.0}
# roll no die, get no result with probability 1

print(weighted_dice_probabilities(1, [1/3, 1/3, 1/3]))
{(0, 0, 1): 0.3333, (0, 1, 0): 0.3333, (1, 0, 0): 0.3333}
# roll 1 die, get each face with probability 1/3

print(weighted_dice_probabilities(2, [1/3, 1/3, 1/3]))
{(0, 0, 2): 0.1111, (0, 1, 1): 0.2222, (0, 2, 0): 0.1111,
 (1, 0, 1): 0.2222, (1, 1, 0): 0.2222, (2, 0, 0): 0.1111}
# roll 2 dice, get each couple of faces with proba 2/9 or each pair with proba 1/9

print(weighted_dice_probabilities(3, [1/3, 1/3, 1/3]))
{(0, 0, 3): 0.037037, (0, 1, 2): 0.111111, (0, 2, 1): 0.111111, (0, 3, 0): 0.037037,
 (1, 0, 2): 0.111111, (1, 1, 1): 0.222222, (1, 2, 0): 0.111111,
 (2, 0, 1): 0.111111, (2, 1, 0): 0.111111, (3, 0, 0): 0.037037}
# roll 3 dice, get each triple with proba 1/27, or each one-face-plus-a-pair with proba 1/9, or get three distinct faces with proba 2/9.

@Stef

Here is the algorithm I followed:

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))


def product(n: int, r: int) -> list:
    return [[i, *p] for i in range(n + 1) for p in product(n, r - 1)] if r else [[]]


def k_values(n: int, r: int) -> list:
    return [i for i in product(n, r) if sum(i) == n]


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

@Stef, what do you think about this algorithm?

def k_values(n: int, r: int):
    length = r

    def filtered_product(_n, _r):
        return (
            [
                v for i in range(_n + 1) for c in filtered_product(_n, _r - 1)
                if (v := [i, *c]) and (sum(v) == _n if len(v) == length else True)
            ]
            if _r else [[]]
        )

    return filtered_product(n, r)

@Stef
How is the code now?

Note: By the way, the function should be f(n, r - 1), not f(n - i, r - 1).

In the meanwhile, I am thankful to you a lot.

def k_values(n, r, stop: int = 0):
    return (
        [
            v for k in k_values(n, r - 1, stop + 1) for i in range(n + 1)
            if len(v := [i, *k]) != r + stop or sum(v) == n
        ]
        if r else [[]]
    )