plot_3Hbefore2H_withp.py

import numpy as np
def get_probs(p1=0.5, p2=0.5):
    """
    Gets the probabilities of 3 consecutive heads before 2,
    2 consecutive before 3 and draw. 
    args:
        p1: P(heads) for the player needing 2 heads.
        p2: P(heads) for the player needing 3 heads.
    """
    start1 = np.array([1,0,0])
    m_3 = get_consecutive_heads_mat(3,p1)
    start2 = np.array([1,0,0,0])
    m_4 = get_consecutive_heads_mat(4,p2)
    # p_n must always be one toss ahead of q_n_minus_1. So, when p_n is at toss 1, q_n_minus_1 
    # should be at 0. When it is at 2, q_n_minus_1 should be at 1 and so on.
    # that is why p_n starts with 1 and goes to 100 while q_n_minus_1 starts with 0 and goes to 99.
    p_n = np.array([np.dot(start1, np.linalg.matrix_power(m_3,n))[0,2]\
                             for n in range(1,101)])
    q_n_minus_1 = np.array([np.dot(start2, np.linalg.matrix_power(m_4,n))[0,2]\
                             for n in range(100)])
    p1wins = sum(q_n_minus_1*(1-p_n))*p2
    p_n = np.array([np.dot(start2, np.linalg.matrix_power(m_4,n))[0,3] for n in range(1,101)])
    q_n_minus_1 = np.array([np.dot(start1, np.linalg.matrix_power(m_3,n))[0,1] for n in range(100)])
    p2wins = sum(q_n_minus_1*(1-p_n))*p1
    return p1wins, p2wins


def plot_probs():
    ps = np.arange(0,1,0.01)
    p_wins = []
    p_losses = []
    p_draws = []
    for p in ps:
        win,loss = get_probs(p2=p)
        p_wins.append(win)
        p_losses.append(loss)
        p_draws.append(1-win-loss)
    plt.plot(ps, p_wins,label='3ConecHeadsWins')
    plt.plot(ps, p_losses,color='orange',label='2ConecHeadsWins')
    plt.plot(ps, p_draws,color='silver',label='draw')
    plt.axhline(0.450595,color='grey')
    plt.axvline(0.768482,color='grey')
    plt.axvline(0.5,color='grey')
    plt.legend()
    plt.show()