Random number with given PDF in Python

Question

I want to generate an integer random number with a probability distribution function given as a list. For example if pdf=[3,2,1] then I like rndWDist(pdf) to return 0,1, and 2, with probabilities of 3/6, 2/6, and 1/6. I wrote my own function for that since I couldn't find it in the random module.

def randintWDist(pdf):
    cdf=[]
    for x in pdf:
        if cdf:
            cdf.append(cdf[-1]+x)
        else:
            cdf.append(x)
    a=random.randint(1,cdf[-1])
    i=0
    while cdf[i]<a:
        i=i+1
    return i

Is there any shorter method to achieve the same result?

Answer 1

This is a duplicate question: Generate random numbers with a given (numerical) distribution

As the first answer there suggest, you might want to use scipy.stats.rv_discrete.

You might use it like that:

from scipy.stats import rv_discrete
numbers = (1,2,3)
distribution = (1./6, 2./6, 3./6)
random_variable = rv_discrete(values=(numbers,distribution))
random_variable.rvs(size=10)

This returns a numpy array with 10 random values.