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]