return a random item in a list, which has different probability.

random_pick.py

# -*- coding: utf-8 -*-

# Reference: Python Cookbook 2nd Ed. p.188

import random


def random_pick(item_list, probabilities):
    """ return a random item in a list, which has different probability.

    Argument:
    item_list (list of item)
    probabilities (list of float. prob. should be 0 <= x <= 1)
    """

    for p in probabilities:
        if not isinstance(p, float):
            raise TypeError("{0} is not float".format(p))
        if p < 0.0 or 1.0 < p:
            raise ValueError("{0} should be between 0.0 to 1.0".format(p))

    x = random.uniform(0, 1)
    cumulative_probability = 0.0

    for item, item_prob in zip(item_list, probabilities):
        cumulative_probability += item_prob
        if x < cumulative_probability:
            break

    return item


if __name__ == '__main__':

    item_list = ["aaa", "bbb", "ccc", "ddd", "eee"]
    probabilities = [0.05, 0.1, 0.15, 0.3, 0.4]

    NUM_ATTEMPTS = 100

    results = {}

    for i in range(1, NUM_ATTEMPTS):
        print("Test #{0}".format(i))
        picked_item = random_pick(item_list, probabilities)
        results[picked_item] = results.setdefault(picked_item, 0) + 1
        print("random pick result: {0}".format(picked_item))

    print("## total result ##")

    for item, count in results.iteritems():
        print("{0}: {1} times".format(item, count))