Spotify Ticket Lottery

spottylotty.py

from __future__ import division
import sys
import math
import random

# http://www.spotify.com/us/jobs/tech/ticket-lottery/

def bc(n, k, f=math.factorial):
    """Binomial coefficient in terms of Factorials
    (n) = ____n!____
    (k)   k!(n - k)!
    """
    assert 0 <= k <= n, "k=%s, n=%s" % (k, n)
    return f(n) / ( f(k) * f(n - k) )

def hyp_dist(k, n, m, N):
    """
    Hypergeometric distribution
                 (m)(N-m)
    P(X = k) = __(k)(n-k)___
                    (N)
                    (n)
    A collection of N items, m white, rest black.
    Formula shows the chance of picking (without replacement)
    k whites when n items are drawn 
    >>> print '%0.9f' % hyp_dist(k=4, n=10, m=5, N=50)
    0.003964583
    >>> print '%0.10f' % hyp_dist(k=5, n=10, m=5, N=50)
    0.0001189375
    """
    return ( bc(m, k) * bc(N-m, n-k) 
           /     bc(N, n)            )

def P_yay_broadway_from_line(line):
    """
    >>> P_yay_broadway_from_line('100 10 2 1')
    '0.1000000000'
    >>> P_yay_broadway_from_line('100 10 2 2')
    '0.1909090909'
    >>> P_yay_broadway_from_line('10 10 5 1')
    '1.0000000000'
    """
    return P_yay_broadway(*map(int, line.split()))

def P_yay_broadway(m, n, t, p, debug=False):
    if debug:
        print m, 'enter lottery'
        print n, 'winners drawn'
        print t, 'tickers per winner'
        print p, 'people in group\n\n'

    w = int(math.ceil(1.0*p/t))
    if debug:
        print 'need %d wins' % w

    c = 0
    if debug:
        print 'range(%s, %s)' % (w, w+p+1)
    for i in range(w, p+1):
        c += hyp_dist(k=i, n=n, m=p, N=m)

    if debug:
        print 'probability = %0.9f' % c
    return '%0.10f' % c


if __name__ == '__main__':
    import sys
    if sys.argv[1:] and sys.argv[1]=='test':
        import doctest
        doctest.testmod()
    else:
        input_line = sys.stdin.readline()
        print P_yay_broadway_from_line(input_line.strip())