Simple Monty Hall simulation

monty_hall.py

import random

N_DOORS = 3
N_REPETITIONS = 1000


def stick(chosen, choice):
    '''
    This ignores the choice and sticks to the chosen door.
    '''
    return chosen


def switch(chosen, choice):
    '''
    This takes the choice of switching.
    '''
    return choice


def monty_hall(chosen, car, doors):
    '''
    This function takes the role of Monty Hall:
    - If the chosen door actually contains the car, it returns
    any of the other doors.
    - If the chosen door is NOT the car, it returns the car.
    '''
    if (chosen == car):
        return random.choice([door for door in doors if door != car])
    return car


def simulate(n_doors, strategy):
    '''
    This function accepts a number of doors and the strategy of sticking or
    switching that the player will take. Then it picks a door at random to
    contain the car, a door at random to be the one the player first chooses,
    and calls monty_hall to simulate Monty's part in the game. It then runs
    the player's strategy (sticking or switching) and returns True if the
    player got the car, and False otherwise.
    '''
    doors = range(n_doors)
    car = random.choice(doors)
    chosen = random.choice(doors)
    remaining = monty_hall(chosen, car, doors)
    final_choice = strategy(chosen, remaining)
    return final_choice == car


def repeated_simulation(strategy):
    '''
    This function runs the simulation for a number of doors as set above,
    a number of times, also set above, and returns the number of wins
    and losses, respectively.
    '''
    runs = [simulate(N_DOORS, strategy) for _ in range(N_REPETITIONS)]
    return runs.count(True), runs.count(False)


def main():
    '''
    This main function runs the simulation for both strategies (sticking and switching)
    and prints the results.
    '''
    sticking = repeated_simulation(stick)
    switching = repeated_simulation(switch)

    print("Sticking results:", sticking)
    print("Switching results:", switching)

if __name__ == '__main__':
    main()