Search code examples
pythonprobabilitydice

Rolling N Non-regular Dice in Constant Time

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?

Solution

  • 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))