Ridder Classic solution 21/02/2020

from scipy import special
import math
import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt

#function to calculate log(n choose k)
def logComb(n,k):
    g = max(k, n-k)
    num = sum([math.log(i) for i in range(g+1, n+1)])
    denom = sum([math.log(i) for i in range(1, n-g+1)])
    return(num-denom)

def oneTerm(x, x1, n, p):
    if (x1+n)%2 != 0 or (x-x1-2*n+200)%4 != 0:
        return (0)

    if (x-x1) < -2*(100-n) or (x-x1) > 2*(100-n):
        return (0)

    comb1 = logComb(n, int((x1+n)/2))
    comb2 = logComb(100-n, int((x-x1-2*n+200)/4))

    if p == 0:
        if ((x1+n)/2) == 0:
            ps = 0
            qs = 0
        else:
            return (0)
    elif p == 1:
        if ((n-x1)/2) == 0:
            ps = 0
            qs = 0
        else:
            return (0)
    else:
        ps = ((x1+n)/2) * math.log(p)
        qs = ((n-x1)/2) * math.log(1-p)

    twos = (n-100) * math.log(2)

    return(math.exp(comb1 + ps + qs + comb2 + twos))

def pWinFnc(x,n,p):
    x1Range = range(-n, n+1, 2)
    return(sum([oneTerm(x, x1, n, p) for x1 in x1Range]))

#determining optimal n
xRange = range(1, 201)

pList = np.linspace(0, 1, 101)
maxPList = []
maxNList = []

for p in pList:
    pWin = [(sum([pWinFnc(x,n,p) for x in xRange]), n) for n in range(0,101)]
    z = max(pWin, key=lambda x:x[0])
    maxPList.append(z[0])
    maxNList.append(z[1])
    print(p)

print(maxPList, maxNList)

fig, ax = plt.subplots()
ax.scatter(pList, maxPList, color='b', s=10)
ax.set(xlabel='p, heads probability of the 1-coin', ylabel='P(Win), Win Probability',
       title='P(Win) vs p')
ax.grid()
ax.set_xticks(np.arange(0, 1.01, step=0.1))
fig.savefig("pwin_vs_p.png")

fig, ax = plt.subplots()
ax.scatter(pList, maxNList, color='r', s=10)
ax.set(xlabel='p, heads probability of the 1-coin', ylabel='Optimal n',
       title='Optimal n vs p')
ax.grid()
ax.set_xticks(np.arange(0, 1.01, step=0.1))
fig.savefig("n_vs_p.png")


#printing example PMFs
p = 0.7
n = 40
pWin = [pWinFnc(x,n,p) for x in range(-200,201)]

fig, ax = plt.subplots()
ax.bar(range(-200, 201), pWin, color='g')
ax.set(xlabel='x, Total score', ylabel='Probability',
       title=f'PMF f(x), n={n}, p={p}')
plt.axvline(x=0, color = "b")
ax.grid()
fig.savefig("PMF.png")

p = 0.01
n = 15
pWin = [pWinFnc(x,n,p) for x in range(-200,201)]

fig, ax = plt.subplots()
ax.bar(range(-200, 201), pWin, color='g')
ax.set(xlabel='x, Total score', ylabel='Probability',
       title=f'PMF f(x), n={n}, p={p}')
plt.axvline(x=0, color = "b")
ax.grid()
fig.savefig("PMF2.png")