Counting and chance in Python: simulating the matching and birthday problems

0. What this note is for

Chapter 1 of Blitzstein and Hwang’s Introduction to Probability ends with a short section of R code: vectors, factorials, random sampling, and two simulations that check results proved earlier in the chapter. This note covers the same ground in Python.

The translation is not line by line. R is built around vectors; Python is built around names, functions, and lists. So the code below follows the habits of CS61A instead:

• every idea gets a small named function with a docstring;
• every docstring carries doctests, examples that Python can check;
• a simulation is a function passed as a value, not a special command.

Everything uses only the standard library (math, random, collections), so it runs in any Python 3 and in the exercises on this page.

1. From vectors to lists

R’s c(3,1,4,1,5,9) becomes a Python list:

>>> v = [3, 1, 4, 1, 5, 9]
>>> sum(v), max(v), min(v), len(v)
(23, 9, 1, 6)

The sequence $1,2,\dots,n$, written 1:n in R, is range(1, n + 1) in Python. A range stops before its second argument, which is why the + 1 is there. Counting down from $m$ to $n$ needs a negative step: range(m, n - 1, -1).

1.1. Positions start at zero

The biggest change is indexing. R numbers positions from 1; Python numbers them from 0. So v[0] is the first entry, and the $i$th entry in the book’s sense is v[i - 1].

>>> v = [3, 1, 4, 1, 5, 9]
>>> v[1]

Picking out the 1st, 3rd, and 5th entries, v[c(1,3,5)] in R, is a list comprehension over positions. Removing the 2nd through 4th entries, v[-(2:4)] in R, is two slices joined together:

>>> [v[i - 1] for i in (1, 3, 5)]
[3, 4, 5]
>>> v[:1] + v[4:]
[3, 5, 9]

A negative index means something else in Python: v[-1] is the last entry, not “everything except the first”.

1.2. Componentwise work is written out

R applies arithmetic to each entry of a vector automatically. In Python a list comprehension says so explicitly. The vector $\left(1,\frac{1}{2^2},\frac{1}{3^2},\dots,\frac{1}{100^2}\right)$ is

>>> squares = [1 / k ** 2 for k in range(1, 101)]
>>> squares[:3]
[1.0, 0.25, 0.1111111111111111]

Powers use **. In Python ^ is bitwise exclusive or, so 2 ^ 3 is 1, not 8. This is one of the easiest mistakes to make when porting R code.

R also recycles a shorter vector: v + 3 adds 3 to every entry. Python lists do not.

>>> v = [3, 1, 4]
>>> v + 3

2. Factorials and binomial coefficients

The standard library has both: math.factorial(n) and math.comb(n, k).

>>> import math
>>> math.factorial(20)
2432902008176640000
>>> math.comb(10, 3)
120

Python integers have no fixed size, so math.factorial(1000) is an exact 2568-digit integer. The limit appears only when we convert to a floating-point number: float(math.factorial(171)) raises OverflowError, the same point at which R returns Inf. For large $n$, work with logarithms. Since $n!=\Gamma(n+1)$, the natural log $\log(n!)$ is math.lgamma(n + 1), and $\log\binom{n}{k}$ is lgamma(n + 1) - lgamma(k + 1) - lgamma(n - k + 1).

Writing $\binom{n}{k}$ ourselves is a good first recursion. The base cases come straight from the definition, and the recursive case is Pascal’s rule.

>>> choose(5, 2)
10

>>> choose(10, 5)
252

>>> choose(4, 0)
1
>>> choose(4, 4)
1

>>> choose(3, 5)
0

def choose(n, k):
    if k < 0 or k > n:
        return 0
    if k == 0 or k == n:
        return 1
    return choose(n - 1, k - 1) + choose(n - 1, k)

3. Random samples

The random module plays the role of R’s sample. Positions in the book are the numbers $1,\dots,n$, so we sample from range(1, n + 1).

random.choices always samples with replacement; the standard library has no weighted sampling without replacement.

These numbers are pseudo-random: a deterministic algorithm produces them. That is useful. Calling random.seed(110) first makes a simulation reproducible, so someone else running your code gets the same numbers.

4. A simulation is a function

R’s replicate(times, expression) evaluates an expression many times. In Python we pass the experiment itself, a function of no arguments, and call it as often as we like:

def replicate(times, experiment):
    """Call EXPERIMENT, a function of no arguments, TIMES times.
    Return the results as a list.

    >>> replicate(3, lambda: 'hi')
    ['hi', 'hi', 'hi']
    """
    return [experiment() for _ in range(times)]

This is the CS61A idea of a higher-order function. replicate knows nothing about cards or birthdays; it only knows how to repeat. The lambda wraps an expression in a function so that it is evaluated later, and anew on each call, rather than once.

5. The matching problem

Shuffle a deck of cards labeled $1,\dots,n$ and turn them over one at a time. A match happens when card $i$ is the $i$th card turned over. The book shows (Example 1.6.4) that the probability of at least one match is

To simulate it we need to count matches in a given deck. In Python, True and False behave as 1 and 0 in arithmetic, just as TRUE and FALSE do in R.

>>> deck = [2, 1, 3, 4]
>>> sum([card == position for position, card in enumerate(deck, start=1)])

>>> count_fixed_points([1, 2, 3])
3

>>> count_fixed_points([2, 1, 3])
1

>>> count_fixed_points([3, 1, 2])
0
>>> count_fixed_points([2, 1])
0

>>> count_fixed_points([])
0

def count_fixed_points(deck):
    return sum(card == position for position, card in enumerate(deck, start=1))

With count_fixed_points written, one shuffle and a full simulation are each a line or two:

import random

def shuffled_matches(n):
    """Shuffle cards 1..N and count the matches."""
    deck = random.sample(range(1, n + 1), n)
    return count_fixed_points(deck)

random.seed(110)
results = replicate(10 ** 4, lambda: shuffled_matches(100))
sum(r >= 1 for r in results) / 10 ** 4     # proportion with a match

6. The birthday problem

With $k$ people and 365 equally likely birthdays (ignoring February 29), the probability that at least two share a birthday is

math.prod computes the product directly, and exact integer arithmetic keeps it accurate until the final division:

>>> k = 23
>>> 1 - math.prod(range(365 - k + 1, 366)) / 365 ** k
0.5072972343239854

As a function with doctests, multiplying one factor at a time:

def birthday_match(k, days=365):
    """Probability that at least two of K people share a birthday.

    >>> round(birthday_match(23), 3)
    0.507
    >>> birthday_match(1)
    0.0
    >>> birthday_match(366)
    1.0
    """
    no_match = 1.0
    for i in range(k):
        no_match *= (days - i) / days
    return 1 - no_match

R’s qbirthday(p) answers the reverse question: how many people are needed for the probability to reach $p$? Python’s standard library has no equivalent, so write it.

>>> people_needed(0.5)
23

>>> people_needed(0.9)
41

>>> people_needed(0.99)
57

>>> people_needed(0.5, days=2)
2

def people_needed(p, days=365):
    k = 1
    while birthday_match(k, days) < p:
        k += 1
    return k

6.1. Simulating instead

To simulate, draw 23 birthdays with replacement and ask how often the most common birthday appears. collections.Counter does what R’s tabulate does:

from collections import Counter

def most_shared(k, days=365):
    """Draw K birthdays; return the size of the largest group sharing one."""
    birthdays = random.choices(range(1, days + 1), k=k)
    return max(Counter(birthdays).values())

random.seed(110)
r = replicate(10 ** 4, lambda: most_shared(23))
sum(x >= 2 for x in r) / 10 ** 4           # at least one shared birthday
sum(x >= 3 for x in r) / 10 ** 4           # at least one triple

The simulation answers questions the formula does not. R’s pbirthday(23, coincident=3) gives about $0.014$ for a triple birthday among 23 people, and qbirthday(0.5, coincident=3) gives 88. Python has no built-in for triples, but most_shared handles them with no extra work.

Unequal birthday probabilities make the formula hard but the simulation easy: pass weights= to random.choices.

7. Comments and docstrings

A # starts a comment in both languages. Python adds the docstring: the string at the top of a function, which says what the function does. With doctests inside, the documentation is also a test. Run

python3 -m doctest simulation.py -v

and Python checks every >>> example in the file. Docstrings for what a function promises, comments for why a line is written the way it is.

References

• Joseph K. Blitzstein and Jessica Hwang, Introduction to Probability, Chapter 1, §1.6 (Example 1.6.4, the matching problem) and §1.8 (R). Course materials: stat110.net.
• CS61A and John DeNero’s Composing Programs, for the style: higher-order functions, doctests, and “What would Python display?”
• Python documentation: random, math, and collections.Counter.