multinomial sample generator

gistfile1.py

import numpy as np

def multinomial_sample(n,p):
    """Return sample variates from multinomial.
    
    In contrast with `numpy`'s multinomial sampler, this will return the actual
    samples, rather than the summed results. In other words, the result of
    `numpy`'s sampler is the same length as `p`, while the result of this
    sampler is of length `n`.
    
    `n` - number of experiments
    `p` - vector of parameters; must sum to 1
    """
    if sum(p) != 1.: raise ValueError, "p must sum to 1"
    uniform_sample = np.random.uniform(size=n)
    p_cum = np.cumsum(p)
    return np.searchsorted(p_cum,uniform_sample,side='right')