table genders problem from stat 61

table-genders.py

from __future__ import division

import random
from collections import defaultdict

def find_pmf(fun, iters=100000):
    counts = defaultdict(int)
    for i in range(iters):
        counts[fun()] += 1

    return dict((key, val / iters) for key, val in counts.iteritems())

def expectation(pmf):
    return sum(key * val for key, val in pmf.iteritems())


def make_table(n):
    arr = [True] * n + [False] * n
    random.shuffle(arr)
    return arr

def is_man_next_to_woman(i, arr):
    if not arr[i]:
        return False
    elif i == 0:
        return not arr[-1] or not arr[1]
    elif i == len(arr) - 1:
        return not arr[-2] or not arr[0]
    else:
        return not arr[i-1] or not arr[i+1]

def count_tablemen(arr):
    return sum(1 if is_man_next_to_woman(i, arr) else 0 for i in range(len(arr)))

def men_at_table(n):
    return count_tablemen(make_table(n))
tabler = lambda n: lambda: men_at_table(n)

def experimental(n):
    return expectation(find_pmf(tabler(n)))

def theoretical(n):
    return 3/2 * n*n / (2*n - 1)

for i in range(1, 11):
    print "%2d %10.3f %10.3f" % (i, experimental(i), theoretical(i))