Calculate the distribution of outcomes when rolling multiple dice.

dice.py

#! /usr/bin/python3


#  This script will calculate the distribution of rolling any number
#  of mixed-sided dice.


#  Example usage:  python3  dice.py  3d6  1d20
#
#  The above would calculate the distribution of rolling four dice
#  all at the same time: 3d6 + 1d20.


#  Data Representation
#
#  dice.py operates on distributions.
#
#  A distribution is modeled as an array, where the index represents
#  the total rolled, and the value represents the count of that total
#  (i.e. how many times that total occurs in the distribution).
#
#  For example, below are the distributions of 1d6, 2d6, and 3d6.
#
#  dist_1d6  =  [  0,  1,  1,  1,  1,  1,  1  ]
#  dist_2d6  =  [  0,  0,  1,  2,  3,  4,  5,  6,  5,  4,  3,  2,  1  ]
#  dist_3d6  =  [  0,  0,  0,  1,  3,  6, 10, 15, 21, 25, 27,
#                             27, 25, 21, 15, 10,  6,  3,  1  ]


#  The roll() function calculates the product of any two
#  distributions.
#
#  roll() operates in O(n*m), where n and m are the lengths of the
#  two arrays that respectively represent each distribution.
#
#  Or perhaps somewhat worse than O(n*m), as with larger n and m (and
#  with larger counts) the multiplication and addition of very large
#  integers will also be involved.


#  Rolling more than two dice is done iteratively, by adding one die
#  at a time to the cumulative distribution.
#
#  For example, dist_4d6 is calculated as follows:
#
#  dist_1d6  =  [0] + [1] * 6
#  dist_2d6  =  roll ( dist_1d6, dist_1d6 )
#  dist_3d6  =  roll ( dist_2d6, dist_1d6 )
#  dist_4d6  =  roll ( dist_3d6, dist_1d6 )
#
#  Note that any two distributions can be combined.  So, for example,
#  dist_4d6 could also be calculated as follows:
#
#  dist_4d6  =  roll ( dist_2d6, dist_2d6 )
#
#  This means that a further optimization is possible, but is not
#  implemented below.
#
#  Consider calculating 1024d6.  At present, dice.py would calculate
#  this as follows:
#
#  dist_1d6  =  [0] + [1] * 6
#  dist_2d6  =  roll ( dist_1d6, dist_1d6 )
#  dist_3d6  =  roll ( dist_2d6, dist_1d6 )
#  dist_4d6  =  roll ( dist_4d6, dist_1d6 )
#  ...
#  dist_1023d6  =  roll ( dist_1022d6, dist_1d6 )
#  dist_1024d6  =  roll ( dist_1023d6, dist_1d6 )
#
#  However, this could be greatly accelerated by doubling each
#  distribution as follows:
#
#  dist_1d6  =  [0] + [1] * 6
#  dist_2d6  =  roll ( dist_1d6, dist_1d6 )
#  dist_4d6  =  roll ( dist_2d6, dist_2d6 )
#  dist_8d6  =  roll ( dist_4d6, dist_4d6 )
#  ...
#  dist_512d6   =  roll ( dist_256d6, dist_256d6 )
#  dist_1024d6  =  roll ( dist_512d6, dist_512d6 )
#
#  Distributions that are not powers of two could be calculated by
#  rolling together the approprate powers of two.  For example:
#
#  dist_768d6  =  roll ( dist_512d6, dist_256d6 )
#
#  Implementing this performance improvement is left as an exercise
#  for the reader.


import  re, sys


def  roll  ( a, b ):    #  ---------------------------------------------  roll

  #  roll() returns the product of distributions a and b.

  dist  =  [ 0 ] * ( len(a) + len(b) - 1 )    #  the new distribution
  for    na  in  range ( 0, len(a) ):
    for  nb  in  range ( 0, len(b) ):
      dist [ na + nb ]  +=  a[na] * b[nb]

  return  dist


def  roll_ndn  ( dist, ndn ):    #  --------------------------------  roll_ndn

  #  roll_ndn() returns the product of dist and ndn.
  #
  #  dist is any pre-existing distribution.
  #
  #  ndn is a string that names a distribution.
  #  for example, ndn could be '2d6' or '4d20'

  try:
    m    =  re .match ( '^(\d+)d(\d+)$', ndn )    #  parse  ndn
    d    =  int ( m.group(1) )    #  how many dice?
    s    =  int ( m.group(2) )    #  how many sides?
    die  =  [0] + [1] * s         #  the distribution for an s-sided die
  except:  print ( 'bad arg  ' + ndn )  ;  sys .exit (1)

  for  j  in  range(0,d):         #  roll die into dist, d times
    dist  =  roll ( dist, die )

  return  dist


def  print_dist ( dist ):    #  ---------------------------------  print_dist
  for  n  in  range ( 0, len(dist) ):
    if  dist[n] > 0:
      print ( ( '%-6d  %d' ) % ( n, dist[n] ) )


def  main  ( argv ):    #  --------------------------------------------  main
  dist  =  [ 1 ]    #  the empty distribution (always zero)
  for  ndn  in  argv:                  #  for each arg in argv
    dist  =  roll_ndn ( dist, ndn )    #  roll in ndn
  print_dist ( dist )


main ( sys .argv [1:] )