Weighted random choice in Python

There are two tiny issues I’d like to address today: first, there is no method in Python’s random module for weighted random choice; second, I haven’t posted anything for too long ;) So, let’s go through a very simple way to implement a function that chooses an element from a list, not uniformly, but using a given weight for each element.

Input: a collection \(C\) of elements and a probability distribution \(p\) over \(C\);
Output: an element chosen at random from \(C\) according to \(p\).

Suppose \(C\) has \(n\) elements numbered from 1 to \(n\), and let us write:

\begin{equation*} p = (p_1, ..., p_n) \end{equation*}

What we have is basically the random.random() function that gives us uniformly at random a floating-point number in \([0, 1]\). Our idea is to partition \([0, 1]\) into \(n\) segments of length \(p_1\) to \(p_n\). This is indeed a partition since \(p\) is a probability distribution and:

\begin{equation*} \sum_{i=1}^n p_i = 1 \end{equation*}

If we draw uniformly at random a point in \([0, 1]\), the probability that it ends up in the \(i^{\mathrm{th}}\) segment is then \(p_i\). Therefore, we can just call random.random(), look where the resulting number falls, and return the corresponding index from the partition. In Python, this becomes:

def weighted_choice(collection, weights):
    assert len(weights) == len(collection)
    assert abs(1. - sum(weights)) < 1e-6
    x = random.random()
    for i, element in enumerate(collection):
        if x <= weights[i]:
            return element
        x -= weights[i]