coin_flip.py

"""
C6 = sum of heads on the 60% coin after N flips
C5 = sum of heads on the 50% coin after N flips


P(C6 > C5 | N flips) >= 0.95 # solve for smallest N

C6 ~ Binomial(N, 0.6)
C5 ~ Binomial(N, 0.5)


define the pmf of d = C6 - C5, p(d). I want to know sum p(d) for 
1) positive d (I guess the coin with more flips is biased)
2) a tie (when I randomly guess, so multiply by 0.5)

"""

from utils.utils import nCr, memoized

@memoized
def m(n, k, p):
    return nCr(n, k) * p ** k * (1 - p) ** (n - k)


def p(d, n, p1, p2):
    return sum( m(n, i, p2)*m(n, d+i, p1) for i in xrange(0, n+1))


def guess_right(n, p1=0.6, p2=0.5):
    running_sum = 0.5*p(0, n, p1, p2)
    d = 1
    while p(d, n, p1, p2) > 0:
        running_sum += p(d, n, p1, p2)
        d += 1
    return running_sum

print guess_right(133)
print guess_right(134)


index = np.arange(0.51, 0.99, step=0.03)
cols = np.arange(5, 140, 5)

results = []
for p1 in index:
    _results = []
    print p1
    for n in cols:
        _results.append(guess_right(n, p1=p1))
    results.append(_results)

df = pd.DataFrame(results, index=index, columns=cols)
sns.heatmap(df[np.arange(5, 140, 5)], cmap='RdYlGn_r', linewidths=0)
plt.title('Probability of guessing correctly')
plt.xlabel("number of flips")
plt.ylabel("P(heads) in biased coin")