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