Dice probability problem

[Antoon Pardon]

Op 2006-04-04, Tomi Lindberg schreef <tomi.lindberg.NO_SPAM at pp.inet.fi.invalid>:
> First, thanks to Antoon and Alexander for replying.
>
> Antoon Pardon wrote:
>
>> It would be better to construct distributions for one
>> die and make a function that can 'add' two distributions
>> together.
>
> As both replies pointed to this direction, I tried to take 
> that route. Here's the unpolished code I came up with. Does 
> it look even remotely sane way to accomplish my goal?

> -- code begins --
>
> # A die with n faces
> D = lambda n: [x+1 for x in range(n)]
>
> # A new die with 6 faces
> d6 = D(6)
>
> # Adds another die to results.
> def add_dice(sums, die):
>      # If first die, all values appear once
>      if not sums:
>          for face in die:
>              sums[face] = 1
>      # Calculating the number of appearances for additional
>      # dice
>      else:
>          new_sums = {}
>          for k in sums.keys():
>              for f in die:
>                  if new_sums.has_key(k+f):
>                      new_sums[k+f] += sums[k]
>                  else:
>                      new_sums[k+f] = sums[k]
>          sums = new_sums
>      return sums
>
> sums = add_dice({}, d6)
> sums = add_dice(sums, d6)
> sums = add_dice(sums, d6)
>
> -- code ends --

IMO you are making things too complicated and not general
enough. Here is my proposal.

----- 

import operator

class Distribution(dict):

  '''A distribution is a dictionary where the keys are dice
     totals and the values are the number of possible ways
     this total can come up '''

  def __add__(self, term):
   
    '''Take two distributions and combine them into one.'''
    
    result = Distribution()
    for k1, v1 in self.iteritems():
      for k2, v2 in term.iteritems():
        k3 = k1 + k2
        v3 = v1 * v2
        try:
          result[k3] += v3
        except KeyError:
          result[k3] = v3
    return result

  def __rmul__(self, num):
    tp = num * [self]
    return reduce(operator.add, tp)

def D(n):

  ''' One die has a distribution where each result has
      one possible way of coming up '''
  return Distribution((i,1) for i in xrange(1,n+1))


sum3d6 = 3 * D(6)
sum2d6p2d4 = 2 * D(6) + 2 * D(4)

-----

-- 
Antoon Pardon