Weighted choice short and simple
PythonNumpyPython Problem Overview
If I have a collection of items in a list. I want to choose from that list according to another list of weights.
For example my collection is ['one', 'two', 'three']
and the weights are [0.2, 0.3, 0.5]
, the I would expect the method to give me 'three' in about half of all draws.
What is the easiest way to do so ?
Python Solutions
Solution 1 - Python
Since [tag:numpy] version 1.7 you can use numpy.random.choice()
:
elements = ['one', 'two', 'three']
weights = [0.2, 0.3, 0.5]
from numpy.random import choice
print(choice(elements, p=weights))
Solution 2 - Python
Since Python 3.6, you can do weighted random choice (with replacement) using random.choices
.
> random.choices(population, weights=None, *, cum_weights=None, k=1)
Example usage:
import random
random.choices(['one', 'two', 'three'], [0.2, 0.3, 0.5], k=10)
# ['three', 'two', 'three', 'three', 'three',
# 'three', 'three', 'two', 'two', 'one']
Solution 3 - Python
This function takes two arguments: A list of weights and a list containing the objects to choose from:
from numpy import cumsum
from numpy.random import rand
def weightedChoice(weights, objects):
"""Return a random item from objects, with the weighting defined by weights
(which must sum to 1)."""
cs = cumsum(weights) #An array of the weights, cumulatively summed.
idx = sum(cs < rand()) #Find the index of the first weight over a random value.
return objects[idx]
It does not use any python loops.
Solution 4 - Python
You can use the multinomial distribution (from numpy) to do what you want. E.g.
elements = ['one', 'two', 'three']
weights = [0.2, 0.3, 0.5]
import numpy as np
indices = np.random.multinomial( 100, weights, 1)
#=> array([[20, 32, 48]]), YMMV
results = [] #A list of the original items, repeated the correct number of times.
for i, count in enumerate(indices[0]):
results.extend( [elements[i]]*count )
So the element in first position came up 20 times, the element in second position came up 32 times, and the element in third position came up 48 times, roughly what you would expect given the weights.
If you're having a hard time wrapping your head around the multinomial distribution, I found the documentation really helpful.
Solution 5 - Python
If you did not want to use numpy
, you can follow the same method with something like this:
from random import random
from itertools import takewhile
def accumulate(iterator):
"""Returns a cumulative sum of the elements.
accumulate([1, 2, 3, 4, 5]) --> 1 3 6 10 15"""
current = 0
for value in iterator:
current += value
yield current
def weightedChoice(weights, objects):
"""Return a random item from objects, with the weighting defined by weights
(which must sum to 1)."""
limit = random()
return objects[sum(takewhile(bool, (value < limit for value in accumulate(weights))))]
We use itertools.takewhile()
to avoid checking values once we reach the point we want to stop, otherwise, this is essentially the same idea as Mischa Obrecht's answer, just without numpy
.
Solution 6 - Python
How about just initializing your list to match your choices with the expected weights. Here I'm make a list of 100 values representing your desired "pull" percentage.
>>> import random
>>> elements = ['one', 'two', 'three']
>>> weights = [0.2, 0.3, 0.5]
>>>
>>> # get "sum" of result list of lists (flattens list)
>>> choices = sum([[element] * int(weight * 100)for element, weight in zip(elements, weights)], [])
>>> random.choice(choices)
three
It's not cumulative, but it looks like it might be what your looking for.
Solution 7 - Python
To build upon Maus' answer, which is great if you want to repeatedly get weighted random values, if you only wanted a single value, you can do this very simply by combining numpy.random.multinomial()
and itertools.compress()
:
from itertools import compress
from numpy.random import multinomial
def weightedChoice(weights, objects):
"""Return a random item from objects, with the weighting defined by weights
(which must sum to 1)."""
return next(compress(objects, multinomial(1, weights, 1)[0]))