Peak balls from a bin (Sampling wihtout Replacement)

Bin3.py

#
# Peak balls from a bin (Sampling wihtout Replacement)
#

import math
import matplotlib.pyplot as plt
import numpy as np

#
# n Choice r
#

def combination (n, r): 
    f = math.factorial
    return f(n) // f(r) // f(n-r)

#
# Probability of the question.
#    

def prob (n, mu):
    n1 = int (n * mu)
    n2 = n - n1
    p1 = combination (n1, 1) * combination (n2, 9) / combination (n, 10)
    p2 = combination (n1, 0) * combination (n2, 10) / combination (n, 10)
    p = p1 + p2
    return p


if __name__ == '__main__':
    print ("prob (100, 0.4)    = ", prob (100, 0.4))
    print ("prob (1000, 0.4)   = ", prob (1000, 0.4))
    print ("prob (10000, 0.4)  = ", prob (10000, 0.4))

    largestY = prob (100000, 0.4)   
    print ("prob (100000, 0.4) = ", largestY)

    x = np.arange (100, 3000, 10)
    y = [prob (x, 0.4) for x in x]

    plt.plot (x, y, label = 'y = prob (x, 0.4)')
    plt.axhline (y = largestY, color='r', linestyle='-', label='y = %f' % largestY)
    plt.legend ()
    plt.show ()