A demonstration of the Monty Hall problem.

monty_hall.py

#!/usr/bin/python3

"""A demonstration of the Monty Hall problem."""

import argparse
import random
import json

# n is the number of times you switch, and also the number of times you don't switch
# accuracy increases as n increases
# n can be overridden with the first command line argument passed, e.g.
# ./monty_hall.py 314
# sets n=314
# ./monty_hall.py
# sets n to what it is set to below:
n = 1000
DOORS = 3  # MUST be >=3
DOORS_LOWER_BOUND = 0
# what to subtract from `DOORS` in order to get the upper bound
doors_subtrahend = 1 - DOORS_LOWER_BOUND
doors_upper_bound = DOORS - doors_subtrahend

def main():
    global n
    args = parse_args()
    n = args.iterations
    switching_probability = car_when_switching_probability()
    stats = simulate(DOORS, n)
    print_output(args, stats, switching_probability)
    return 0

def car_when_switching_probability(door_count: int=DOORS, return_ratio: bool=True):
    """Calculates the probability of getting a car when you switch. Assumes there is 1 car.
    
    door_count: How many doors there are.
    return_ratio: Whether you want a float returned, or a 2-tuple of integers representing the ratio
                  (in the format of (numerator, denominator))"""
    # probability of getting a goat as your first choice
    # represented as a 2-tuple of the ratio
    first_choice_goat_prob = (door_count - 1, door_count)
    # probability that you will switch to a car, given that you got a goat the first time
    switch_to_car_prob = (1, door_count - 2)
    prob = []
    # multiply the fractions (because it's the probability of them both happening)
    for i in range(2):  # 2 = length of the above tuples
        prob.append(first_choice_goat_prob[i] * switch_to_car_prob[i])
    
    if return_ratio:
        # not going to simplify the fraction because it will always result in either
        # odd/even
        # or even/odd
        # so there's less chance of the ratio not being coprime
        # for all the smaller numbers of doors, the ratio is already coprime
        # idk whether or not there are any non-coprime ratios and i don't think it's worth it to check
        return tuple(prob)
    else:
        return prob[0]/prob[1]

class ObjectWithAttrs:
    """A generic class for storing attributes. init function takes kwargs which are the attributes."""
    def __init__(self, **kwargs):
        for key, value in kwargs.items():
            setattr(self, key, value)

def simulate(door_count: int, _n: int):
    """Simulates the Monty Hall problem ``n`` times for switching, and ``n`` times again for not switching.
    
    Returns an object with a ``switching`` and ``not_switching`` attribute, both of which contain
    integers representing how many times the player won a car."""
    switching = 0
    not_switching = 0
    for switch in (True, False):  # iterate twice: once for switching, once for not switching
        for j in range(_n):  # now simulate Monty Hall n times
            result = monty_hall(door_count, switch)
            if result:
                if switch:
                    switching += 1
                else:
                    not_switching += 1
    return ObjectWithAttrs(switching=switching, not_switching=not_switching)

def monty_hall(door_count: int, switch: bool):
    """Complete one simulation of the Monty Hall problem.
    
    Returns a ``bool`` representing whether or not the player won a car."""
    
    car_door = randomly_choose_door()
    first_choice = randomly_choose_door()
    opened_door = open_a_door(door_count, car_door, first_choice)
    final_choice = switch_doors(door_count, first_choice, opened_door) if switch else first_choice
    
    return final_choice == car_door

def parse_args():
    """Handles argument parsing with the ``argparse`` module."""
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("iterations", action="store", default=n, type=int, nargs="?",
        help="How many times the Monty Hall problem will be simulated "
        "for both switching and not switching doors respectively. (default: %(default)s)")
    parser.add_argument("-j", "--json", action="store_true", dest="json",
        help="Output JSON instead of human-friendly output.")
    return parser.parse_args()

def randomly_choose_door(lower_bound: int=DOORS_LOWER_BOUND, upper_bound: int=doors_upper_bound):
    return random.randint(lower_bound, upper_bound)

def switch_doors(door_count: int, first_choice: int, opened_door: int):
    """Chooses a random door to switch to which is not the one which has been opened nor the first choice door.
    Returns an int representing which door has been switched to."""
    # 2 representing the first choice door and the opened door, both of which can't be chosen now
    remaining_doors = door_count - 2
    choice = random.randint(0, remaining_doors - doors_subtrahend)
    cannot_choose = [opened_door, first_choice]  # doors which cannot be switched to
    # sorting in order to avoid a situation where
    # first_choice = opened_door - 1
    # choice = first_choice
    # therefore
    # choice = opened_door due to adding 1 for ``first_choice`` but have already processed ``opened_door``
    cannot_choose.sort()
    # account for said doors
    # can be imagined as, say, [0, 1, 2, opened_door, 3, 4, first_choice]
    # becoming [0, 1, 2, 3, 4, 5, 6], where opened_door is 3 and first_choice is 6
    # so 3 and 4 have become 4 and 5
    for door in cannot_choose:
        if choice >= door:
            choice += 1
    # if choice is over the upper door bound, set it to the lower door bound
    # this can be imagined as a circular set like {... 1, 2, 0, 1, 2 ...}
    # so instead of being 3, it would be 0. or instead of being 4, it would be 1
    if choice > doors_upper_bound:
        choice -= door_count
    final_choice = choice
    return final_choice

def open_a_door(door_count: int, car_door: int, first_choice: int):
    """Returns a goat door to open which is neither the first choice door nor the car door."""
    opened_door = None
    # choosing a door to open is not random as which door is opened doesn't affect the odds
    # so long as it's a non-prize door which is opened
    # and obviously your first choice door can't be opened either as that defeats the purpose
    for door in range(door_count):
        if door not in (car_door, first_choice):
            opened_door = door
            break
    
    if opened_door is None:
        raise ValueError("All doors were either player's first chosen door or a car door. "
            "This probably means that DOORS < 3. DOORS must be >=3. "
            "You can fix this by editing the constant assignment at the top of the script.")
    
    return opened_door

def print_output(args, stats, switching_probability):
    """Handles printing the output."""
    if args.json:
        output = {
            "iterations": n,
            "doors": DOORS,
            "switching": stats.switching,
            "not_switching": stats.not_switching
        }
        print(json.dumps(output, indent=4))
    else:
        print(f"Cars attained when switching: {stats.switching}/{n} "
        f"(mathematically expected is ~{switching_probability[0]}/{switching_probability[1]})")
        print(f"Cars attained when not switching: {stats.not_switching}/{n} (mathematically expected is ~1/{DOORS})")

main()