Weighted choices

[Peter Otten]

Jason Friedman wrote:

> choices = dict()
> choices["apple"] = 10
> choices["pear"] = 20
> choices["banana"] = 15
> choices["orange"] = 25
> choices["kiwi"] = 30
> 
> I want to pick sets of fruit, three in a set, where the chance of
> selecting a given fruit is proportional to its weight.  In the example
> above, pears should appear twice as often as apples and kiwis should
> appear twice as often as bananas.
> 
> I see this solution
> (http://stackoverflow.com/questions/3679694/a-weighted-version-of-random-
choice)
> but it is for single choices at a time.  random.sample() does not
> provide a "weight" argument.
> 
> I can picture how to write this with many, many lines of code but I
> have a feeling there is an elegant solution.

Let's assume for a moment that you always have low integral weights. Given 

data = []
for weight, value in choices.items():
   data.extend([value]*weight)

picking values with

N = 3
sample = [random.choice(data) for i in range(N)]

does not give te same probabilities as

sample = random.sample(data, N)

Which one doe you want? If it's the former, i. e. picking an item does not 
affect the item's probability, just adapt Raymond Hettinger's

def weighted_choice(choices):
    values, weights = zip(*choices)
    total = 0
    cum_weights = []
    for w in weights:
        total += w
        cum_weights.append(total)
    x = random() * total
    i = bisect(cum_weights, x)
    return values[i]

to (untested)

def weighted_choice_iter(choices):
    values, weights = zip(*choices)
    total = 0
    cum_weights = []
    for w in weights:
        total += w
        cum_weights.append(total)
    while True:
        x = random() * total
        i = bisect(cum_weights, x)
        yield values[i]

def weighted_choices(choices, N):
    return list(itertools.islice(weighted_choice_iter(choices), N))

print weighted_choices(choices, 3)