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.
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 dice each with sides, the number of equally likely ways to roll a score is given by
so that the probability of rolling a score is .
If we now consider keeping only of the dice, then we can group the desired outcomes with score by:
- the smallest value among the largest dice that we keep,
- the number of kept dice that are strictly greater than , and
- the number of discarded dice that are strictly less than .
This yields the following summation:
At this point, we’re done… except that if the number 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
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
- Goetgheluck, P., Computing Binomial Coefficients, The American Mathematical Monthly, 94(4) April 1987, p. 360-365 [JSTOR]