Generalized Birthday Problem

generalized_birthday_problem.py

import sys
import math

n = int(sys.argv[1])    # n is the number of people
M = int(sys.argv[2])    # M is the number of matching birthdays
m = int(sys.argv[3])    # m is the number of days in the year

def logFact(n):
    """
    Compute log(n!)
    """
    return logFallingFact(n, n)

def logFallingFact(n, k):
    """
    Compute the log of the "falling factorial"
    
    n * (n - 1) * ... * (n - k)
    """
    f = 0
    for j in range(n, n - k, -1):
        f += math.log(j)
    return f

def T(n, m, M):
    """
    Compute the probability that there is at least one M-tuple of matching
    birthdays for n people and m day-long years.
    """
    return 1.0 - P(n, m, m, M - 1)

def P(n, mOrig, mCur, k):
    """
    Compute the probability that there are at most k-tuples of matching
    birthdays (i.e. no r-tuples where r > k) for n people and mOrig day-long
    years.

    There can be at most n // k such k-tuples, so we sum the probabilities
    for each possible case (0, 1, ..., n // k).
    """
    if k == 1:
        plog = logFallingFact(mCur, n)
        plog -= math.log(mOrig) * n
        return math.exp(plog)
    p = 0
    for i in range(0, n // k + 1):
        p += Q(n, mOrig, mCur, k, i)
    return p

def Q(n, mOrig, mCur, k, i):
    """
    Compute the probability that there are i k-tuples of matching birthdays
    for n people and mOrig day-long years.

    Let f(n, k) be the "falling factorial" n * (n - 1) * ... (n - k + 1).
    Let f(n) be the factorial f(n, n).

    Then the number of ways to pick i k-tuples, and label them in sorted
    order of the index of their first element, is:

    f(n, k * i) / (f(k) ** i * f(i))

    We then assign i of the mCur remaining dates to those k-tuples, in
    sorted order of the index of their first element.  (This insistence on
    sorting is to prevent double-counting valid arrangements.)

    Finally, we need to figure out what the probability is that the
    remaining (n - k * i) elements have at most (k - 1)-tuples of
    matching birthdays.  We've used i days, so there are now (mCur - i)
    days left to assign.  Fortunately, we already have a function for this
    above ;)
    """
    plog = logFallingFact(n, k * i) + logFallingFact(mCur, i)
    plog -= logFact(i) + logFact(k) * i + math.log(mOrig) * k * i
    return math.exp(plog) * P(n - k * i, mOrig, mCur - i, k - 1)

print(T(n, m, M))