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.