from pylab import * # library for making plots
from decimal import * # library for exact arithmetic with large numbers
# make a plot of the probability mass function y(x)
def pmfplot(x, y):
bar(x, y, width=0.5, align='center')
xticks(x)
ylim(0, 1.1 * max(y))
show()
# calculate the binomial coefficients (n choose k) for k from 0 to n
def binom(n):
B = [0]*(n + 1)
B[0] = 1
for k in range(1, n + 1):
B[k] = B[k - 1] * (n - k + 1) / k
return B
# the probability that you collect all of n coupons in d days
def coupons_prob(n, d):
B = binom(n)
sum = 0
for i in range(1, n):
sum = sum + (-1)**(i + 1) * B[i] * (n - i)**d
return 1.0 - float(Decimal(sum) / Decimal(n**d))
# plot the probability of collecting all coupons within 1, ..., d days
def plot_coupons_prob(n, d):
days = range(0, d + 1)
plot(days, map(lambda day: coupons_prob(n, day), days), 'x')
ylim([0.0, 1.0]) # sets the boundary values of the y axis
show() # displays the plot
# find the first day on which all coupons are collected with probability at least p
def day_of_coupon_collection(n, p):
# do a binary search on d starting at d = n
d_l = n-1
d_r = n
p_r = coupons_prob(n, d_r)
while(p_r < p):
d_l = d_r
p_l = p_r
d_r = 2 * d_l
p_r = coupons_prob(n, d_r)
while(d_r - d_l > 1):
d_m = (d_l + d_r)/2
if(coupons_prob(n, d_m) < p):
d_l = d_m
else:
d_r = d_m
return d_r
# plot the day on which the probability of coupon collection becomes p for up to n coupons
def plot_coupons_day(n, p):
coupons = range(1, n + 1)
plot(coupons, map(lambda nc: day_of_coupon_collection(nc, p), coupons), 'x')
show() # displays the plot
# plot the approximating function n log(n / log(1 / (1 - p)))
def plot_coupons_day_appx(n, p):
coupons = range(1, n + 1)
log1p = math.log(1.0 / (1.0 - p))
plot(coupons, map(lambda c: c * math.log(c / log1p), coupons))
show() # displays the plot