gistfile1.py

#!/usr/bin/env python

'''
Id:          "$Id$"
Description: 
Test:        
'''

"""
Given a list of probabilities (floats), and a list of values (anything), write a function/class/whatever:

go(probs, values)

that returns the values with the given probabilities.
For example, you'd expect calling:

go([0.3, 0.7], ['hello', 'world'])

100 times would give (roughly, but probably not exactly) 30 'hello's and 70 'world's, in some order.

Bear in mind that you don't know in advance how many times 'go' is going to be called, so it's basically a random value generator.

There are many ways to write this, many interfaces -- go could be a function, a class, whatever you want.
Try to do it in the most Pythonic way possible, and write some unit tests.

"""

import random
import unittest2
#import itertools
#from nose.tools import eq_, raises
#from collections import defaultdict
#from contextlib import contextmanager
try:
    from qz.core.test.qzunittest import qzUnitTest
except ImportError:
    pass

from exceptions import Exception

class ProbTestInitException(Exception):
    pass

class ProbTest(object):
    
    def __init__(self, probs, values):
        """
        We expect two equal-sized lists of probabilities and values.
        """
        if len(probs) != len(values):
            raise ProbTestInitException
        if sum(probs) != 1.0:
            raise ProbTestInitException

        self.gen = None

        combined = zip(probs, values)
        
        self.generator_list = []
        for prob, val in combined:
            self.generator_list += ([val,]*int(prob*10))  # i.e. create a list of 'foo' and put them on the list

    def _gen(self, gen_vals):
        for i in range(len(gen_vals)):
            yield gen_vals[i]
            
    def go(self):
        if not self.gen:
            self.gen = self._gen(self.generator_list)
        return self.gen.next()

class TestGo(unittest2.TestCase):
    
    def test_pass(self):
        pass
        
    def test_failure_unless_both_len_of_probs_and_values_are_equal(self):
        probs = [0.1,]
        values = ['foo', 'bar,']
        with self.assertRaises(ProbTestInitException):
            ProbTest(probs, values)
        
    def test_failure_unless_probs_values_sum_to_one(self):
        probs = [0.1, 0.4, 0.2,]
        values = ['foo', 'bar', 'baz', ]
        with self.assertRaises(ProbTestInitException):
            ProbTest(probs, values)

    def test_successful_init(self):
        probs = [0.4, 0.6,]
        values = ['foo', 'bar', ]
        pt = ProbTest(probs, values)
        self.assertEqual(isinstance(pt, ProbTest), True)

    def test_always_expected_value(self):
        probs = [1.0,]
        values = ['foo', ]
        pt = ProbTest(probs, values)
        self.assertEqual(pt.go(), 'foo')
        
    def test_values(self):
        probs = [0.4, 0.6,]
        values = ['foo', 'bar', ]
        pt = ProbTest(probs, values)
        print pt.go()

            
def main():    
    #qzUnitTest(headless=True, suppressOutput=False)
    
    probs = [0.1, 0.2, 0.4, 0.3,]
    values = ['one', 'two', 'three', 'four', ]
    pt = ProbTest(probs, values)
    print pt.generator_list
    print len(pt.generator_list)
    count_dict = {}
    for i in range(1000):
        try:
            key = pt.go()
            count_dict[key] = count_dict.setdefault(key, 0) + 1
        except StopIteration:
            break
    print count_dict

if __name__ == "__main__":
    main()