Conditional probability in Python: two children, Monty Hall, and updating on evidence

0. Where this picks up

This note continues Counting and chance in Python, which rebuilt Chapter 1’s R section. Chapter 2 of Blitzstein and Hwang’s Introduction to Probability is about conditional probability, and its R section (§2.10) checks two famous results by simulation: the two-child problem and the Monty Hall problem. The same habits apply: small named functions, doctests, and functions passed around as values. The standard library is enough, and replicate from the previous note is not needed here.

1. Conditional probability as a frequency

The frequentist reading of conditional probability says: repeat an experiment $n$ times, count the repetitions where $B$ happened ($n_B$) and where both $A$ and $B$ happened ($n_{AB}$). For large $n$,

We only look at the runs where $B$ happened, and ask how often $A$ happened among those.

1.1. Simulating families

Let each family have two children, each a girl (1) or a boy (2) with equal probability, independently. R builds two vectors, one per child, and Python can do the same. zip then pairs them up, one tuple per family:

import random

random.seed(110)
n = 10 ** 5
child1 = random.choices([1, 2], k=n)     # elder child of each family
child2 = random.choices([1, 2], k=n)     # younger child
families = list(zip(child1, child2))

families[0] might be (1, 2): an elder girl and a younger boy. Keeping the two children of a family in one tuple means we can never accidentally pair the elder child of one family with the younger child of another.

1.2. No elementwise comparison

R’s child1 == 1 compares every entry of a vector with 1 and returns a vector of TRUE and FALSE. Python’s == compares whole values, and a list is not equal to a number.

>>> child1 = [1, 2, 1]
>>> child1 == 1

So Python spells out the loop with a comprehension. R’s & and | become Python’s and and or, applied inside it:

>>> pairs = [(1, 1), (1, 2), (1, 1)]
>>> [c1 == 1 and c2 == 1 for c1, c2 in pairs]
[True, False, True]

>>> pairs = [(1, 1), (1, 2), (2, 1), (1, 1)]
>>> sum([c1 == 1 and c2 == 1 for c1, c2 in pairs])

1.3. Events are predicates

Rather than writing a new count for every question, describe each event as a predicate: a function that takes an outcome and returns True or False. Then one function estimates any conditional probability. This is the CS61A idea of passing functions as arguments.

>>> families = [(1, 1), (1, 2), (2, 1), (2, 2)]
>>> both_girls = lambda f: f == (1, 1)
>>> conditional_frequency(families, both_girls, lambda f: f[0] == 1)
0.5

>>> families = [(1, 1), (1, 2), (2, 1), (2, 2)]
>>> both_girls = lambda f: f == (1, 1)
>>> round(conditional_frequency(families, both_girls, lambda f: 1 in f), 3)
0.333

>>> families = [(1, 1), (1, 2), (2, 1), (2, 2)]
>>> both_girls = lambda f: f == (1, 1)
>>> conditional_frequency(families, both_girls, lambda f: False) is None
True

>>> rolls = [1, 2, 3, 4, 5, 6]
>>> conditional_frequency(rolls, lambda r: r == 6, lambda r: r % 2 == 0)
0.3333333333333333

def conditional_frequency(outcomes, a, b):
    n_b = sum(1 for o in outcomes if b(o))
    n_ab = sum(1 for o in outcomes if a(o) and b(o))
    return None if n_b == 0 else n_ab / n_b

With conditional_frequency written, the book’s two questions are two lines:

both_girls = lambda f: f == (1, 1)
elder_girl = lambda f: f[0] == 1
at_least_one_girl = lambda f: 1 in f

conditional_frequency(families, both_girls, elder_girl)
conditional_frequency(families, both_girls, at_least_one_girl)

2. The Monty Hall problem

A car is behind one of three doors, goats behind the other two. You pick a door. Monty, who knows where the car is, opens one of the other doors to show a goat, and offers to let you switch. Should you?

2.1. Never switching

Assume you always pick door 1. Then the never-switch strategy wins exactly when the car is behind door 1, and we do not even need to simulate Monty:

random.seed(110)
cardoor = random.choices([1, 2, 3], k=n)
sum(car == 1 for car in cardoor) / n

One run with seed 110 gave $0.3313$, close to $1/3$.

2.2. Modeling Monty

To simulate switching we need Monty’s choice. When your door hides a goat, Monty has exactly one door he can open. When your door hides the car, he has two, and in the game he picks at random. For code that is easy to test, the next exercise makes a fixed choice instead, always the lowest-numbered door allowed. Randomness can be added back around it later.

>>> monty_opens(1, 2)
3
>>> monty_opens(2, 1)
3
>>> monty_opens(3, 1)
2

>>> monty_opens(1, 1)
2
>>> monty_opens(2, 2)
1
>>> monty_opens(3, 3)
1

def monty_opens(chosen, car):
    return min(door for door in (1, 2, 3) if door != chosen and door != car)

Once Monty has opened a door, switching means taking the one door that is neither yours nor his. Doors 1, 2, and 3 add up to 6, so that door is 6 - chosen - opened.

To check that answer, simulate the whole game, with Monty choosing at random when he has a choice:

def play_switching():
    """Pick door 1, watch Monty open a goat door, switch; return True on a win."""
    car = random.randint(1, 3)
    chosen = 1
    opened = random.choice([d for d in (1, 2, 3) if d != chosen and d != car])
    return 6 - chosen - opened == car

random.seed(110)
sum(play_switching() for _ in range(n)) / n

2.3. Playing the game

The book ends with an interactive monty() function. In Python, input plays the role of R’s scan: it prints a prompt and returns what the player types, always as a string, so int(...) converts the door number. An f-string replaces R’s paste, and startswith replaces substr(reply, 1, 1).

import random

def switch_door(chosen, opened):
    """The one door that is neither CHOSEN nor OPENED."""
    return 6 - chosen - opened

def monty():
    """Play one round of the Monty Hall game at the terminal."""
    doors = [1, 2, 3]
    car = random.choice(doors)                       # where the car is
    print("Monty Hall says 'Pick a door, any door!'")
    chosen = int(input('Your door (1, 2, or 3): '))
    # Monty opens a door that is neither yours nor the car's.
    opened = random.choice([d for d in doors if d != chosen and d != car])
    print(f'Monty opens door {opened}!')
    reply = input('Would you like to switch (y/n)? ')
    if reply.strip().lower().startswith('y'):        # "yes", "yep", "Yeah!"
        chosen = switch_door(chosen, opened)
    print('You won!' if chosen == car else 'You lost!')

Save this as monty.py and run python3 -i monty.py, then call monty() as often as you like. The exercises on this page cannot run it, because they have no keyboard input.

Notice the shape of the program. Every decision that can be tested without a person at the keyboard (switch_door, and monty_opens in the exercise) is a pure function. Only the outer monty talks to the player. CS61A calls this separating computation from input and output, and it is why the earlier exercises can check the logic with doctests.

3. Posteriors become priors

The end of Chapter 2 (Figure 2.7) describes how evidence arrives in steps. Before any evidence, $P(A)$ is the prior. After learning $B$, the posterior $P(A\mid B)$ comes from Bayes’ rule. That posterior is the prior for the next step, and learning $C$ gives $P(A\mid B,C)$.

>>> bayes_update(Fraction(1, 4), 1, Fraction(2, 3))
Fraction(1, 3)

>>> bayes_update(Fraction(1, 3), 1, Fraction(1, 2))
Fraction(1, 2)

>>> bayes_update(Fraction(1, 100), 1, 0)
Fraction(1, 1)

>>> bayes_update(Fraction(1, 5), Fraction(1, 2), Fraction(1, 2))
Fraction(1, 5)

def bayes_update(prior, p_evidence_if_a, p_evidence_if_not_a):
    numerator = p_evidence_if_a * prior
    return numerator / (numerator + p_evidence_if_not_a * (1 - prior))

Updating one piece of evidence at a time is then a chain of calls, each output feeding the next input. Take the two-child family again, with $A$ = both children are girls, so $P(A)=1/4$.

>>> monday = bayes_update(Fraction(1, 4), 1, Fraction(2, 3))
>>> monday
Fraction(1, 3)

>>> bayes_update(monday, 1, Fraction(1, 2))
Fraction(1, 2)

fractions.Fraction keeps the arithmetic exact, so $1/3$ stays $1/3$ rather than becoming $0.333\ldots$ with a rounding error.

References

• Joseph K. Blitzstein and Jessica Hwang, Introduction to Probability, Chapter 2: Example 2.2.5 (two children), the Monty Hall problem, Figure 2.7 (sequential updating), and §2.10 (R). Course materials: stat110.net.
• CS61A and John DeNero’s Composing Programs, for higher-order functions, doctests, and separating logic from input and output.
• Python documentation: random, fractions, and input.