Skip to content

Instantly share code, notes, and snippets.

@vyznev

vyznev/dice_roll.py

Last active May 30, 2021
Embed
What would you like to do?
A simple bare-bones dice probability calculator framework, compatible with both Python 2 and Python 3
def dice_roll(die, count = 1, select = None):
"""Generate all possible results of rolling `die` `count` times, sorting
the results (according to the order of the sides on the die) and selecting
the first `select` elements of it.
The yielded results are tuples of the form `(roll, prob)`, where `roll` is a
sorted tuple of `select` values and `prob` is the probability of the result.
The first argument can be either a custom die, i.e. a tuple of `(side, prob)`
pairs, where `prob` is the probability of rolling `side` on the die, or just
a simple integer, which will be passed to `make_simple_die`.
Keyword arguments:
die -- a custom die or an integer
count -- the number of dice to roll (default 1)
select -- the number of results to select (set equal to count if omitted)
"""
# cannot select more dice than there are in the pool
if select is None or select > count:
select = count
# for convienience, allow simple dice to be given as plain numbers
die = make_simple_die(die) if isinstance(die, int) else tuple(die)
if len(die) == 1:
# base case: a one-sided die has only one possible result
yield ((die[0][0],) * select, die[0][1]**count)
elif len(die) > 1:
# split off the first side of the die, normalize the rest
side, p_side = die[0]
rest = tuple((side, prob / (1-p_side)) for side, prob in die[1:])
p_sum = 0 # probability of rolling this side less than select times
for i in range(0, select):
# probability of rolling this side exactly i times
p_i = binomial(count, i) * p_side**i * (1-p_side)**(count-i)
p_sum += p_i
# recursively generate combinations
for roll, p_roll in dice_roll(rest, count-i, select-i):
yield ((side,) * i + roll, p_i * p_roll)
# final case: all selected dice (and possibly more) roll this side
yield ((side,) * select, 1-p_sum)
_factorials = [1]
def binomial(n, k):
"""Helper function to efficiently compute the binomial coefficient."""
while len(_factorials) <= n:
_factorials.append(_factorials[-1] * len(_factorials))
return _factorials[n] / _factorials[k] / _factorials[n-k]
def make_simple_die(n):
"""Generate a simple n-sided die with sides listed in decreasing order."""
return tuple((i, 1.0/n) for i in range(n, 0, -1))
def explode(die, count=2):
"""Make an "exploding die" where the first (=highest) side is rerolled up to
count times.
"""
die = make_simple_die(die) if isinstance(die, int) else tuple(die)
exploded = die
for i in range(count):
top, p_top = exploded[0]
exploded = tuple((side + top, prob * p_top) for side, prob in die) + exploded[1:]
return exploded
def sum_roll(die, count = 1, select = None, ascending=False):
"""Convenience function to sum the results of `dice_roll()`. Takes the same
parameters as `dice_roll()`, returns a list of `(sum, prob)` pairs sorted in
descending order by sum (and thus suitable for use as a new custom die). The
optional parameter `ascending=True` can be used to change the sort order.
"""
from collections import defaultdict
summary = defaultdict(float)
for roll, prob in dice_roll(die, count, select):
summary[sum(roll)] += prob
return tuple(sorted(summary.items(), reverse = not ascending))
@vyznev

This comment has been minimized.

Copy link
Owner Author

@vyznev vyznev commented Mar 21, 2020

This code is released into the public domain under the Creative Commons CC0 Public Domain Dedication. Feel free to do whatever you want with it. Credit is appreciated but not required.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment