The distribution of “roll and keep” dice

Introduction

Some role-playing games involve a “roll and keep” dice mechanic, where you roll some number of dice, but only “keep” a specified number of them with the largest values, where the result of the roll is the sum of the kept dice. For example, rolling five dice and keeping the largest three (sometimes denoted 5d6k3) could yield a score between 3 and 18, inclusive. What is the probability of each possible score?

The motivation for this post is (1) to capture my notes on the general solution to this problem, and (2) to describe some potential optimizations in implementation to handle very large instances of the problem, such as this Hacker Rank version of the problem, where we may be rolling as many as 10,000 dice.

The solution

To start, let’s simplify the problem by just rolling and keeping all of the dice without discarding any, since the order statistics are what make this problem complicated. Given n dice each with d sides, the number r_{n,d}(s) of equally likely ways to roll a score s is given by

r_{0,0}(s) = 1

r_{n,d}(0) = 0

r_{n,d}(s) = \sum\limits_{k=0}^{\lfloor\frac{s-n}{d}\rfloor} (-1)^k {n \choose k} {{s-k d-1} \choose {n-1}}

so that the probability of rolling a score s is r_{n,d}(s)/d^n.

If we now consider keeping only m \leq n of the dice, then we can group the r_{n,d,m}(s) desired outcomes with score s by:

  • the smallest value a among the m largest dice that we keep,
  • the number b of kept dice that are strictly greater than a, and
  • the number c of discarded dice that are strictly less than a.

This yields the following summation:

r_{n,d,m}(s) = \sum\limits_{a=1}^d \sum\limits_{b=0}^{m-1} \sum\limits_{c=0}^{n-m} {n \choose {b,c}} (a-1)^c r_{b,d-a}(s-a m)

Implementation details

At this point, we’re done… except that if the number n of dice rolled is large, then a straightforward implementation of this nested summation will be pretty slow, involving calculation of a large number of very large multinomial coefficients.

My first thought was to speed up calculation of those very large coefficients via their prime factorization using Legendre’s formula. The paper by Goetgheluck linked at the end of this post describes a more optimized approach, but the following simple Python implementation is already significantly faster than the usual iterative algorithm for sufficiently large inputs (the implementation of primes(n) is left as an exercise for the reader, or a future post):

def binomial(n, k):
    """Large binomial coefficient n choose k."""
    if k < 0 or k > n:
        return 0
    c = 1
    for p in primes(n):
        c = c * p ** (v_fact(n, p) - v_fact(k, p) - v_fact(n - k, p))
    return c

def v_fact(n, p):
    """Largest power of prime p dividing n factorial."""
    v = 0
    while n > 0:
        n = n // p
        v = v + n
    return v

(It’s worth noting that the Hacker Rank problem actually only asks for the number of ways to roll a given total modulo a prime, suggesting that a similar but even more effective optimization using Lucas’s theorem is probably intended.)

But we can do much better than this, by observing that in the “usual” iterative algorithm, all of the intermediate products inside the loop are also binomial coefficients… and we need all of them for this problem. That is, we can rewrite the summation as

r_{n,d,m}(s) = \sum\limits_{a=1}^d \sum\limits_{b=0}^{m-1} {n \choose b} r_{b,d-a}(s-a m) \sum\limits_{c=0}^{n-m} {n-b \choose c} (a-1)^c

and then “embed” the usual iterative calculation of the binomial coefficients inside the nested summations. The following implementation eliminates all but a single explicit calculation of a binomial coefficient:

def rolls_keep(s, n, d, m):
    """Ways to roll sum s with n d-sided dice keeping m largest."""
    result = 0
    for a in range(1, d + 1):
        choose_b = 1
        for b in range(0, m):
            sum_c = 0
            choose_c = 1
            pow_c = 1
            for c in range(0, n - m + 1):
                sum_c = sum_c + choose_c * pow_c
                choose_c = choose_c * (n - b - c) // (c + 1)
                pow_c = pow_c * (a - 1)
            result = result + choose_b * rolls(s - a * m, b, d - a) * sum_c
            choose_b = choose_b * (n - b) // (b + 1)
    return result

def rolls(s, n, d):
    """Ways to roll sum s with n d-sided dice."""
    if s == 0 and n == 0:
        return 1
    if d == 0:
        return 0
    result = 0
    choose_k = 1
    for k in range(0, (s - n) // d + 1):
        result = result + (-1) ** k * choose_k * binomial(s - k * d - 1, n - 1)
        choose_k = choose_k * (n - k) // (k + 1)
    return result

Reference:

  1. Goetgheluck, P., Computing Binomial Coefficients, The American Mathematical Monthly, 94(4) April 1987, p. 360-365 [JSTOR]