Empirical Monty Hall

mh.py

import random


def stay_strategy(objects):
    first = random.randint(0, len(objects)-1)
    return objects[first]

def switch_strategy(objects):
    first = random.randint(0, len(objects)-1)
    goat = random.choice([i for i,x in enumerate(objects) if x == 'G' and i != first])
    switch = random.choice([i for i, x in enumerate(objects) if i != first and i != goat])
    return objects[switch]

def immutable_shuffle(objects):
    return random.sample(objects, len(objects))


def gen_pop(size, num_cars=1):
    l = ['G' for i in xrange(size-num_cars)]
    return l + ['C' for i in xrange(num_cars)]


Iterations = 10000
PopSize = 3
NumCars = 1

stay_wins = 0
switch_wins = 0
for i in xrange(Iterations):
    pop = gen_pop(PopSize, NumCars)

    if stay_strategy(pop) == 'C':
        stay_wins += 1
    if switch_strategy(pop) == 'C':
        switch_wins += 1

print 'Stay', stay_wins/float(Iterations)
print 'Switch', switch_wins/float(Iterations)