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:
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:
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]
Discussion
There are no comments yet. Feel free to leave a reply using the form below.