Random number with given PDF in Python
I want to create an integer random number with a probability distribution function specified as a list. For example, if pdf = [3,2,1] then I like rndWDist (PDF) to return 0,1 and 2, with probabilities 3/6, 2/6 and 1/6. I wrote my own function for myself as 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 a shorter method to achieve the same result?
source to share
This is a duplicate question: Generate random numbers with a given (numeric) distribution
As you can see from the first answer, you can use scipy.stats.rv_discrete
.
You can use it like this:
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)
Returns a numpy array with 10 random values.
source to share
Given the format of your input, you can do:
def randint_with_dist(pdf):
choices = []
for index, value in enumerate(pdf):
choices.extend(index for _ in range(value))
return random.choice(choices)
Since the same list will be used every time the same list is passed pdf
, you might consider caching the list for efficiency (at the expense of space):
def randint_with_dist(pdf, choices={}):
pdf = tuple(pdf)
if pdf not in choices:
choices[pdf] = []
for index, value in enumerate(pdf):
choices[pdf].extend(index for _ in range(value))
return random.choice(choices[pdf])
source to share
Using numpy (version 1.7 or newer) you can also use np.random.choice :
In [27]: import numpy as np
In [28]: distribution = (1./6, 2./6, 3./6)
In [29]: np.random.choice(np.arange(len(distribution)), p=distribution)
Out[29]: 0
In [30]: np.random.choice(np.arange(len(distribution)), p=distribution, size=10)
Out[30]: array([2, 1, 1, 2, 2, 0, 1, 0, 1, 0])
source to share