Dibujo rechoncho de urna

Quiero realizar un sorteo aleatorio relativamente simple en números, pero no puedo encontrar una buena manera de expresslo. Creo que la mejor manera es describirlo como un dibujo de una urna sin reemplazo. Tengo una urna con k colores, y n_k bolas de todos los colores. Quiero dibujar m bolas y saber cuántas bolas de cada color tengo.

Mi bash actual es

np.bincount(np.random.permutation(np.repeat(np.arange(k), n_k))[:m], minlength=k)

Aquí, n_k es una matriz de longitud k con los conteos de las bolas.

Parece que es equivalente a np.bincount(np.random.choice(k, m, n_k / n_k.sum(), minlength=k)

que es un poco mejor, pero aun así no es genial.

Lo que quieres es una implementación de la distribución hipergeométrica multivariable . No conozco a nadie que esté adormecido o con miedo, pero puede que ya exista en algún lugar.

Puede implementarlo utilizando llamadas repetidas a numpy.random.hypergeometric . Si eso será más eficiente que su implementación, depende de cuántos colores hay y cuántas bolas de cada color.

Por ejemplo, aquí hay un guión que imprime el resultado del dibujo de una urna que contiene tres colores (rojo, verde y azul):

 from __future__ import print_function import numpy as np nred = 12 ngreen = 4 nblue = 18 m = 15 red = np.random.hypergeometric(nred, ngreen + nblue, m) green = np.random.hypergeometric(ngreen, nblue, m - red) blue = m - (red + green) print("red: %2i" % red) print("green: %2i" % green) print("blue: %2i" % blue) 

Salida de muestra:

 red: 6 green: 1 blue: 8 

La siguiente función generaliza que al elegir m bolas dadas una matriz de colors contienen el número de cada color:

 def sample(m, colors): """ Parameters ---------- m : number balls to draw from the urn colors : one-dimensional array of number balls of each color in the urn Returns ------- One-dimensional array with the same length as `colors` containing the number of balls of each color in a random sample. """ remaining = np.cumsum(colors[::-1])[::-1] result = np.zeros(len(colors), dtype=np.int) for i in range(len(colors)-1): if m < 1: break result[i] = np.random.hypergeometric(colors[i], remaining[i+1], m) m -= result[i] result[-1] = m return result 

Por ejemplo,

 >>> sample(10, [2, 4, 8, 16]) array([2, 3, 1, 4]) 

Lo siguiente debería funcionar:

 def make_sampling_arr(n_k): out = [ x for s in [ [i] * n_k[i] for i in range(len(n_k)) ] for x in s ] return out np.random.choice(make_sampling_arr(n_k), m)