Python doesn't have any weighted sampling functionality built in (NumPy/SciPy does), but for a really simple case like this, it's pretty easy:
import itertools
import random
probabilities = [0.3, 0.2, 0.5]
totals = list(itertools.accumulate(probabilities))
def sample():
n = random.uniform(0, totals[-1])
for i, total in enumerate(totals):
if n <= total:
return i
If you don't have Python 3.2+, you don't have the accumulate
function; you can fake it with an inefficient one-liner if the list really is this short:
totals = [sum(probabilities[:i+1]) for i in range(len(probabilities))]
… or you can write an explicit loop, or an ugly reduce
call, or copy the equivalent Python function from the docs.
Also, note that random.uniform(0, totals[-1])
is just a more complicated way of writing random.random()
if you can be sure that your numbers add up to 1.0.
A quick way to test this:
>>> samples = [sample() for _ in range(100000)]
>>> samples.count(0)
29878
>>> samples.count(1)
19908
>>> samples.count(2)
50214
Those are pretty close to 30%, 20%, and 50% of 100000, respectively.
与恶龙缠斗过久,自身亦成为恶龙;凝视深渊过久,深渊将回以凝视…