montecarlo()

montecarlo.py

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

from decimal import Decimal
import random


def montecarlo(choices):
    
    """
    It's easier for me to explain this with an example.
    
        >>> choices = {
        ...     0.2: 1
        ...     0.3: 2
        ...     0.5: 3
        ... }
    
    There's a 20% probability of producing 1, 30% of producing 2, and 50% of
    producing 3.
    
        >>> montecarlo(choices)
        3
    
    The probabilities don't have to add up to 1, either; they could have been
    4, 6 and 10 and it would have been the same.
    """
    
    cumulative_probabilities = {Decimal(0): None}
    probability_sum = Decimal(0)
    sorted_choices = sorted(choices.items(), key=lambda pair: pair[0])
    
    for probability, outcome in sorted_choices:
        if isinstance(probability, (Decimal, int, long)):
            probability_sum += probability
        
        elif isinstance(probability, float):
            # This will typically round up a bit, but floating points lose
            # some precision anyway.
            probability_sum += Decimal(str(probability))
        
        elif isinstance(probability, basestring):
            probability_sum += Decimal(probability)
        
        cumulative_probabilities[probability_sum] = outcome
    
    rand_number = Decimal(str(random.random())) * probability_sum
    prev = 0
    sorted_probs = sorted(
        cumulative_probabilities.items(), key=lambda pair: pair[0])
    
    for cum_prob, outcome in sorted_probs:
        if prev < rand_number <= cum_prob:
            prev, result = cum_prob, outcome
    return result