binomial_coin_toss.py

import sys
import math
import numpy as np
from matplotlib import pyplot as plt

# https://en.wikipedia.org/wiki/Binomial_distribution
# k successes, n trials, p probability of success
def binomial_pmf(k, n, p):
    nk = math.factorial(n) / (math.factorial(k) * math.factorial(n-k))
    return nk * p**k * (1.0-p)**(n-k)

def graph_coin_tosses(n):
    x = range(0, int(n)+1)
    y = [binomial_pmf(i, int(n), 0.5) for i in x]
    
    plt.scatter(x, y, color="blue")
    plt.xlabel('Number of successes')
    plt.title(''.join(['Binomial distribution for ', n, ' coin tosses']))
    plt.show()

if __name__ == "__main__":
    graph_coin_tosses(sys.argv[1])