Birthday Problem Variation

birthday.py

#! /usr/bin/python

# See http://pballew.blogspot.com/2011/01/back-of-envelope-answers-to-hard.html

import numpy as np

def transition_matrix(slots):
  return np.matrix([
      ([0] * i) +
      [(i+1.0)/slots, 1-(i+1.0)/slots][:slots-i] +
      ([0] * (slots-i-2))
      for i in range(slots)])

def prob_full(people, slots):
  m = np.matrix(transition_matrix(slots))
  p = m ** (people-1)
  return p[0, slots-1]

if __name__ == '__main__':
  print "Four grades, seven people: %f" % prob_full(7, 4)
  print "365 days, 1000 people: %e" % prob_full(1000, 365)
  print "365 days, 2287 people: %f" % prob_full(2287, 365)