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, ..., 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 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, _ in enumerate(seq): if x <= weights[i]: return elmt x -= weights[i]