Probability-weighted drawing.

weighted_draw.py

from __future__ import division
import random

def weighted_draw(items, target):
    """Draw item from iterable based on probability distribution of items.
    Args
        items: tuple of (item: float of probability) pairs where the.
               sum probability has been normalized to 1
        target: float, random number between 0 and 1 used to determine
                the item to draw from items
    """
    cum_prob = 0

    for item, prob in items:
        cum_prob += prob
        if cum_prob > target: break

    return item

def test_weighted_draw(prob_dict, seed, trials):
    """Test weighted_draw() with repeated trials.
    Args
        prob_dict: dict of (item: float of probability) pairs where the.
                   sum probability has been normalized to one.
        seed: int of seed to feed into random.seed()
        trials: int of number of trials to run
    """
    random.seed(seed)

    items = tuple(prob_dict.items())
    count_dict = dict([(key, 0) for key in prob_dict])

    for _ in xrange(trials):
        draw = weighted_draw(items, random.random())
        count_dict[draw] += 1

    for key in sorted(count_dict):
        stat = 'expected: ' + '{:.1%}'.format(prob_dict[key]) + ' / '
        stat += 'observed: ' + '{:.1%}'.format(count_dict[key] / trials)
        print key, stat

if __name__ == '__main__':
    animals = {'zebra': 0.5, 'elephant': 0.25, 'lion': 0.25}

    for exp in (4, 5, 6):
        print '\n', 10 ** exp, 'Trials'
        test_weighted_draw(animals, 123456789, 10 ** exp)