A collapsing die simulator

collapsing_die.py

import random

def collapse_die(dice_count, dice_size):
    '''
    Collapse on a 1, rerolling remaining dice.
    All 1s => 0.
    '''
    # Error checking.
    if dice_count < 1:
        raise Exception("dice_count must be a positive integer")
    if dice_size < 1:
        raise Exception("dice_size must be a positive integer")

    dice_left = dice_count
    while True:
        rolls = []
        for die in xrange(dice_left):
            rolls.append(random.randint(1, dice_size))
        rolls.sort()

        # Count the 1s.
        ones_rolled = len(filter(lambda r: r == 1, rolls))

        # If all 1s, the result is 0.
        if ones_rolled == dice_left:
            return 0

        # Remove all ones.
        dice_left -= ones_rolled
        for one in xrange(ones_rolled):
            rolls.pop(0)

        # Reroll a die for each one rolled.
        if ones_rolled >= len(rolls) / 2.0:
            # More 1s than non 1s, or equal numbers.
            rolls = []
        elif ones_rolled:
            # Fewer than half are ones. 
            for one in xrange(ones_rolled):
                # Prioritize the highest rolls.
                rolls.pop()
        else:
            return sum(rolls)

def roll_trials(repetitions, dice_count, dice_size):
    results = {}
    for n in xrange(repetitions):
        result = collapse_die(dice_count, dice_size)
        if result in results:
            results[result] += 1
        else:
            results[result] = 1
    formatted_results = []
    for k, v in results.iteritems():
        formatted_results.append([k, v])
    formatted_results.sort(cmp=lambda x, y: x[0] < y[0])
    formatted_results = map(lambda x: [x[0], round(float(x[1]) / repetitions, 4)],
                            formatted_results)
    for result in formatted_results:
        print str(result[0]) + " => " + str((100 * result[1])) + "%" 

roll_trials(500000, 5, 6)