Quiero generar un número aleatorio entero con una función de distribución de probabilidad dada como una lista. Por ejemplo, si pdf = [3,2,1], me gusta que rndWDist (pdf) devuelva 0,1 y 2, con probabilidades de 3/6, 2/6 y 1/6. Escribí mi propia función para eso ya que no pude encontrarla en el módulo aleatorio.

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

¿Hay algún método más corto para lograr el mismo resultado?

4
Hooman 3 sep. 2014 a las 14:59

3 respuestas

La mejor respuesta

Esta es una pregunta duplicada: Genera números aleatorios con una distribución (numérica) dada

Como sugiere la primera respuesta, puede usar scipy.stats.rv_discrete.

Puede usarlo así:

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)

Esto devuelve una matriz numpy con 10 valores aleatorios.

4
Community 23 may. 2017 a las 11:57

Dado el formato de su entrada, puede hacer:

def randint_with_dist(pdf):
    choices = []
    for index, value in enumerate(pdf):
        choices.extend(index for _ in range(value))
    return random.choice(choices)

Como se usará la misma lista cada vez que se pase la misma pdf, podría considerar almacenar en caché la lista para una mayor eficiencia (a costa del espacio):

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])
1
jonrsharpe 4 sep. 2014 a las 14:23

Usando numpy (versión 1.7 o más reciente), también puede usar 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])
0
unutbu 3 sep. 2014 a las 12:04