Say I have a non-regular dice defined by probabilities in a list that add up to one, e.g [0.1, 0.3, 0.4, 0.2]
. I can use the following code to simulate rolling that dice n
times:
import random from collections import Counter def roll(dist, n): sides = [i + 1 for i in range(len(dist))] draws = random.choices(population=sides, weights=dist, k=n) return Counter(draws) print(roll([0.1, 0.3, 0.4, 0.2], 10000000))
Counter({3: 4000343, 2: 2998523, 4: 2000309, 1: 1000825})
However, for large n
, the code gets quite slow, as choices
iterates n
times.
Is there an algorithm which can simulate the dice rolls for any n
in constant time?
Advertisement
Answer
Numpy has this built in as numpy.random.Generator.multinomial:
import numpy as np from collections import Counter def roll(dist, n): rng = np.random.default_rng() results = rng.multinomial(n, dist) return Counter(enumerate(results, start=1))