Newer
Older
donut-odds / odds.py
#!/usr/bin/python

import sys
import random

debug = False

def fact(n):
  if n == 0: return 1
  return n * fact(n - 1)

success_odds_by_dice = {}

sides = 6

def get_odds(num_dice, indent = 0):
  global sides, debug, success_odds_by_dice
  def iprint(val, *args):
    if debug:
      print (indent * " ") + str(val).format(*args)

  if num_dice == 0: return 1

  iprint("-- get_odds({}) --", num_dice)

  if num_dice in success_odds_by_dice:
    iprint("-- odds: {} --", success_odds_by_dice[num_dice])
    return success_odds_by_dice[num_dice]

  total_odds_success = 0

  num_rolls = pow(sides, num_dice)
  iprint("num rolls with {} dice: {:,}", num_dice, num_rolls)
  for i in xrange(1, num_dice + 1):
    remaining_dice = num_dice - i
    num_rolls_with_i_sixes = pow(sides - 1, remaining_dice) * fact(num_dice) / (fact(remaining_dice) * fact(i))
    for j in xrange(i + 1, num_dice + 1):
      num_rolls_with_i_sixes 
    iprint(" num rolls of {} dice with {} {}'s: {:,}", num_dice, i, sides, num_rolls_with_i_sixes)
    odds_i_sixes = num_rolls_with_i_sixes / float(num_rolls)
    iprint("  odds of that happening: {}", odds_i_sixes)
    total_odds_success += get_odds(num_dice - i, indent + 3) * odds_i_sixes

  success_odds_by_dice[num_dice] = total_odds_success

  iprint("-- odds: {} --", total_odds_success)
  return total_odds_success

def get_monte_carlo_odds(iterations, num_dice):
  global debug
  successes = 0
  for i in xrange(0, iterations):
    dice_left = num_dice
    while dice_left > 0:
      num_sixes = 0
      for j in xrange(0, dice_left):
        if random.randint(1, sides) == sides:
          num_sixes += 1
      if num_sixes == 0:
        dice_left = 0
      else:
        dice_left -= num_sixes
        if dice_left == 0:
          successes += 1
  if debug:
    print "-- monte carlo: {:,} / {:,} --".format(successes, iterations)
  return successes / float(iterations)

def find_flag(flag):
  if flag in sys.argv:
    sys.argv.remove(flag)
    return True
  return False

debug = find_flag("debug")
simple_output = find_flag("simple")
do_monte_carlo = not simple_output and not find_flag("no_mc")

dice = int(sys.argv[1])

if len(sys.argv) > 2:
  sides = int(sys.argv[2])

monte_carlo_iters = 100000
if len(sys.argv) > 3:
  monte_carlo_iters = int(sys.argv[3])

if do_monte_carlo:
  monte_carlo_odds = get_monte_carlo_odds(monte_carlo_iters, dice)

odds = get_odds(dice)

if debug:
  # extra line to separate debug output
  print
if simple_output:
  print odds
else:
  odds_str = "odds of getting donuts with {} {}-sided dice: {:.2}%".format(dice, sides, odds * 100)
  if do_monte_carlo:
    odds_str += ", monte carlo with {:,} iterations gives: {:.2}%".format(monte_carlo_iters, monte_carlo_odds * 100)
  print odds_str