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:
p=(p1,...,pn)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 p1 to
pn. This is indeed a partition since p is a probability
distribution and:
i=1∑npi=1If we draw uniformly at random a point in [0,1], the probability that
it ends up in the ith segment is then pi.
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
Feel free to post a comment by e-mail using the form below. Your e-mail address will not be disclosed.