birthday distribution PMF for Python 3 (no dependencies)

bday.py

from functools import reduce
from operator import mul

def prod(iterable, initial=1):
    # Return the product of an iterable.
    return reduce(mul, iterable, initial)

def fall(n, k):
    # Compute the falling factorial $n^{(k)}$.
    return prod(range(n, n - k, -1))

def fact(n):
    # Compute the factorial $n!$.
    return fall(n, n)

def binom(n, k):
    # Compute the binomial coefficient $\tbinom{n}{k}$.
    return fall(n, k) // fact(k)

def multinom(n, ks):
    # Compute the multinomial coefficient $\tbinom{n}{k_i}$,
    # where $k_i$ are entries in the iterable 'ks'.
    if not ks:
        return 1
    k, *ks = ks
    return binom(n, k) * multinom(n - k, ks)

def codes(n, k):
    # Compute all codes of size $k$ in $A_n$
    for i in range(n // k, 0, -1):
        for code in subcodes(n - i*k, k - 1):
            code.append(i)
            yield tuple(code)

def subcodes(n, k):
    if k == 0:
        if n == 0:
            yield []
        return
    for i in range(n // k, -1, -1):
        for code in subcodes(n - i*k, k - 1):
            code.append(i)
            yield code

def groups(code):
    g = []
    for i, alpha in enumerate(code, start=1):
        g.extend([i] * alpha)
    return g

def coeff(n, code):
    g = groups(code)
    return multinom(n, g) // prod(fact(x) for x in code)

def q(b, n, k):
    return sum(coeff(n, code) * fall(b, sum(code))
               for code in codes(n, k))

def pmf(b, n, k):
    return q(b, n, k) / b**n