"Monty Hall" problem, generalized to any number of doors, with the host choosing a variable-sized subset of the doors that were not chosen by the player

montyhall.py

#!/usr/bin/env python3


## Generalized Monty Hall problem for N doors, with the host opening some subset
## of the doors that were not choson.

import random
import argparse

def main():
    parser = argparse.ArgumentParser()
    parser.add_argument('--doors', type=int, default=3, help='Total number of doors')
    parser.add_argument('--opens', type=int, default=1, help='Number of doors that Monty opens')
    parser.add_argument('--trials', type=int, default=1000000, help='Number of trials')
    args = parser.parse_args()

    assert args.opens < args.doors-1, f'NUmber of doors opened must me less than the total number of doors minus 1'

    wins = {'stay': 0, 'switch': 0}
    losses = {'stay': 0, 'switch': 0}
    words = ['staying', 'switching']

    for _ in range(args.trials):

        doors = list(range(args.doors))

        firstchoice = random.choice(doors)

        cardoor = random.choice(doors)

        doors.remove(firstchoice)
        if firstchoice != cardoor:
            doors.remove(cardoor)

        # Monty opens some that don't have cars and are not the one that we picked
        for j in range(args.opens):
            doors.pop()

        if firstchoice != cardoor:
            doors.append(cardoor)

        decision = random.choice(['stay', 'switch'])
        if decision == 'switch':
            nextchoice = random.choice(doors)
        else:
            nextchoice = firstchoice

        if nextchoice == cardoor:
            wins[decision] += 1
        else:
            losses[decision] += 1

    stay_switch_ratio = wins["stay"]/max(wins["switch"], 1)
    switch_stay_ratio = wins["switch"]/max(wins["stay"], 1)
    pre_prob = 1 / args.doors
    possible_opens = args.doors - 2
    open_ratio = args.opens / possible_opens # ratio of actual opens to potential opens
    post_prob = ((args.doors - 1) / (args.doors - args.opens - 1)) / args.doors

    print(f'doors={args.doors}, opens={args.opens}, trials={args.trials}')
    print(f'{wins=}, {losses=}')
    print(f'all wins={wins["stay"]+wins["switch"]}, all losses={losses["stay"]+losses["switch"]}')
    print(f'all stays={wins["stay"]+losses["stay"]}, all switches={wins["switch"]+losses["switch"]}')
    print(f'switch/stay win ratio={switch_stay_ratio} ({stay_switch_ratio})')
    best = 0 if wins['stay'] > wins['switch'] else 1
    worst = 1 - best
    print(f'{words[best]} is better than {words[worst]}')
    print(f'pre-reveal probability = 1/{args.doors} ({pre_prob})')
    print(f'{open_ratio=}')
    print(f'post-reveal probability = ({((args.doors - 1) / (args.doors - args.opens - 1))}/{args.doors}) ({post_prob})')
    print(f'post/pre = {post_prob/pre_prob}')


if __name__ == '__main__':
    main()