[Tutor] Weighted Random Choice - Anyone have an efficient algorithm?
Steven D'Aprano
steve at pearwood.info
Thu Dec 23 13:30:50 CET 2010
Modulok wrote:
> Does anyone know of an efficient way of doing a weighted random
> choice? (I don't even know what algorithms like this would be called.)
If you google for "python weighted random choice" you will find a number
of hits.
> Preferably, something that doesn't grow exponentially with the number
> of elements in the list, or the size of their respective values.
Here's one method that is linear on the number of elements:
def weighted_choice(weights):
total = sum(weights)
p = random.uniform(0, total) # random float between 0 and total.
assert 0.0 <= p <= total
# Do a linear search for the right index value. If you have a
# huge number of weights, a binary search may be faster.
running_total = 0
for i, weight in enumerate(weights):
running_total += weight
if p <= running_total:
return i
And tested:
>>> import random
>>> weights = [5, 20, 75]
>>> counts = {0:0, 1:0, 2:0}
>>> for i in xrange(1000000):
... i = weighted_choice(weights)
... counts[i] += 1
...
>>> counts
{0: 50252, 1: 199997, 2: 749751}
>>> [n*1e6/100 for n in weights] # expected values
[50000.0, 200000.0, 750000.0]
As you can see, the results are very close to what should be expected,
and of course the difference can be chalked up to random chance.
--
Steven
More information about the Tutor
mailing list