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 which we number from 1 to n, and let us write (p = (p[1], ..., p[n]). What we have is basically the random.random() function that gives us uniformly at random a float in [0, 1]. Our idea here is to partition [0, 1] into n segments of length p[1] to p[n] (this is indeed a correct partition since the sum of p[i]'s is one). If one draws uniformly at random a point in [0, 1], the probability that it ends up in the i-th segment is then p[i]. Therefore, one can just call random.random(), look where the result ends in [0, 1], and return the corresponding index from the partition. In Python, we can write this as follows:

def weighted_choice(seq, weights):
    assert len(weights) == len(seq)
    assert abs(1. - sum(weights)) < 1e-6

    x = random.random()
    for i, elmt in enumerate(seq):
        if x <= weights[i]:
            return elmt
        x -= weights[i]