Brooklyn 99 themed Monty Hall Problem

montyhall.py

# Simulate the monty hall probelm
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
#import tqdm # gives you a progress bar using tqdm.trange(n)


def monty_hall(switch = None):
    # get some doors
    door = np.array([0,1,2])

    # put car behind one door
    car = np.random.choice(door,1)[0]

    # contenstant picks 1 door
    pickA = np.random.choice(door,1)[0]

    # monty must not open the car or contestant door
    # if contestant picked that car, monty has 2 doors
    # if the contestant picked a goat, monty has 1 door
    avail_doors_to_open = np.delete(door,(car,pickA))
    monty_choice = np.random.choice(avail_doors_to_open,1)[0]
    
    # what is the door the contestant can switch to
    pickB = np.delete(door,(pickA,monty_choice))[0]
  
    # will he switch / yes or no
    if switch is None:
        switch = np.random.choice([True,False]) # pick between 0: stay 1: switch

    if switch:
        contenstant_door = pickB
    else:
        contenstant_door = pickA
    
    gamestate = np.zeros(shape=(4,))
    gamestate[0] = car
    gamestate[1] = monty_choice
    gamestate[2] = pickA
    gamestate[3] = contenstant_door
  
    # wins if final contestant door is the same as the car door
    return switch, gamestate, contenstant_door==car




n = 10000
switch = []
win = []
gms = []
for i in range(n):
    sw, gm, wn  = monty_hall()
    switch.append(sw)
    win.append(wn)
    gms.append(gm)

switch = np.array(switch).astype(bool) # switched doors?
win = np.array(win).astype(bool) # wins/losses
gms = np.array(gms) # list of gamestates



a = switch & win
b = win & (~switch)
mpl.style.use('bmh')
plt.rc('font',family='sans-serif')
plt.figure(figsize=(5,4),facecolor='w')
plt.axhline(1/3, 0,1,color='0.25',ls='--',lw=2)
plt.axhline(2/3, 0,1,color='0.25',ls='--',lw=2)
#plt.axhline(1/2, 0,1,color='0.5',ls='--')
plt.plot(a.cumsum()/Nswitch,'-',label ='Dr. Kevin Cozner',color='navy',lw=3)
plt.plot(b.cumsum()/Nstay,'-',label = 'Capt. Raymond Holt',color='lightsalmon',lw=3)
#plt.plot(win.cumsum()[::1]/len(win),'k-',label = 'Random switching (aka Jake Peralta)')
plt.xticks([1000,5000,10000],['1000','5000','10000'],fontsize=12)
plt.yticks([0,1/3,2/3,1],['0','1/3','2/3','1'],fontsize=12)
plt.ylabel('Fraction wins',fontsize=13)
plt.xlabel('Trial',fontsize=13)
plt.title('Monty Hall (10000 x)',fontsize=14,fontweight='bold')
plt.legend(fontsize=14)
ax = plt.gca()
ax.set_facecolor('w')
plt.tight_layout()
plt.show()
plt.savefig('monty_hall_simulation.png')