Coupon Collector problem with unequal probabilities

generalized-coupon.py

"""
A generalized coupon collector problem.

An experiment with n possible outcomes is repeated until each of the possible
outcomes is observed at least once. Each outcome i has a fixed non-zero
probability p[i], and each trial is independent of the previous trials.
What is the expected number of trials?

For example, what is the expected number of times that two dice must be rolled
until all of the possible totals (from 2 to 12) have appeared at least once?
"""

from fractions import Fraction

def coupon(p):
    p = list(map(Fraction, p))
    sum_p = sum(p)
    p = [x / sum_p for x in p]
    n = len(p)
    N = 2 ** n
    a = [0] + [1] * (N - 1)
    pow2 = [2**i for i in range(n)]
    for i in range(1, N):
        sum_p = 0
        for j in range(n):
            k = i ^ pow2[j]
            if k < i:
                a[i] += p[j] * a[k]
                sum_p += p[j]
        a[i] /= sum_p
    return a[N - 1]

e = coupon([1, 2, 3, 4, 5, 6, 5, 4, 3, 2, 1])
print('Exact value: %s' % e)
print('Approximate value: %.5f' % e)