Probability DP, in depth: forward distributions, backward expectations, cycles and exact answers

Two shapes: forward probability and backward expectation

Every probability DP is one of two shapes, and confusing them is the most common source of wrong answers.

Forward probability DP pushes probability mass from the start state outward. P[t][s] is the probability of being in state s after t steps. You initialise P[0][start] = 1 and, for every state and every outcome, add P[t][s] * p(outcome) to the successor. The invariant is that each layer sums to 1, which is also your cheapest bug check. Forward DP answers questions such as the distribution of a sum, the probability of being absorbed by a given step, or the chance a process ever visits a state within a horizon.

Backward expectation DP pulls values from the goal back towards the start. E[s] is the expected remaining cost from state s. Terminal states have E = 0 and every other state satisfies E[s] = cost(s) + sum over s' of p(s, s') E[s']. You evaluate from the goal backwards because each state depends on its successors. Backward DP answers expected steps, expected score, expected cost and, once you add a max, optimal policies.

Pick the direction whose recurrence is linear in the quantity you want. A forward distribution over positions does not give expected turns unless you also track time and sum t * P(absorbed exactly at t), which is far more work.

Designing a Markov state

The state must satisfy the Markov property: given the state, the future is independent of how you got there. This is the whole design problem. If tomorrow's transition depends on something your state does not record, such as how many times a bonus has already been used, the recurrence is silently wrong even though every individual line of code looks right.

Write the transition function first, as a pure function from state and random outcome to next state; every variable it reads belongs in the state. If the state count is too large, look for symmetry. In the coupon collector, the state is not which faces you have seen (64 subsets) but how many (7 values), because every unseen face is equally likely and the identity of the faces never affects a transition.

Worked example: the distribution of dice sums

Start with the forward shape. What is the probability that three fair dice sum to exactly 10, and that ten dice sum to at least 40? Layer t is a distribution over partial sums; each new die convolves it with the uniform distribution on 1 to 6.

from fractions import Fraction

def dice_sum_distribution(n_dice, faces=6):
    dist = {0: Fraction(1)}
    for _ in range(n_dice):
        nxt = {}
        for total, prob in dist.items():
            for f in range(1, faces + 1):
                nxt[total + f] = nxt.get(total + f, 0) + prob / faces
        dist = nxt
        assert sum(dist.values()) == 1          # layer invariant
    return dist

print(dice_sum_distribution(3)[10])                       # 1/8
d = dice_sum_distribution(10)
print(float(sum(p for s, p in d.items() if s >= 40)))     # 0.20497...

Three dice give 27 of 216 outcomes summing to 10, which is exactly 1/8. Ten dice reach 40 or more about 20.5% of the time, at a cost of O(n × S × faces) instead of 6¹⁰ enumerated outcomes. The assert catches any transition that drops or double-counts an outcome on the first layer where it happens.

Worked example: the coupon collector and self-loops

Now the backward shape. How many rolls of a fair die, on average, until every face has appeared? Let k be the number of distinct faces seen. From state k one roll costs 1; with probability k/6 you repeat a face and stay, otherwise you advance. That gives E[k] = 1 + (k/6) E[k] + ((6-k)/6) E[k+1], and the self-loop means E[k] appears on both sides. Do not iterate; rearrange: E[k] = (1 + ((6-k)/6) E[k+1]) / (1 - k/6), which simplifies to E[k] = 6/(6-k) + E[k+1].

from fractions import Fraction as F

E = [F(0)] * 7                      # E[6] = 0: all faces seen
for k in range(5, -1, -1):
    E[k] = (1 + F(6 - k, 6) * E[k + 1]) / (1 - F(k, 6))
print(E[0], float(E[0]))            # 147/10 14.7

The table reads 14.7, 13.7, 12.5, 11, 9 and 6 for k = 0 to 5: after five faces you still expect six more rolls to see the last one. The lesson generalises. A self-loop is a cycle of length one, and any state whose only cycle is a self-loop can be solved by moving its own term to the left-hand side, which keeps the evaluation order acyclic.

Cycles: linear systems and value iteration

Longer cycles need more. Consider a board with squares 0 to 9. Each turn you roll a fair three-sided die (1, 2 or 3), you need an exact roll to land on 9 (overshoots leave you in place), and square 7 is a chute back to square 1. The chute creates a cycle 1 → ... → 7 → 1, so no evaluation order makes every successor ready before its predecessor. You have two correct options.

Solve the linear system. Each non-terminal state gives one equation, E[s] - sum p(s, t) E[t] = 1, with terminal values fixed at 0. Nine unknowns, nine equations, Gaussian elimination in O(n³). This is exact and is the right tool up to a few thousand states. Value iteration. Start from E = 0 everywhere and repeatedly apply the recurrence until the largest change falls below a tolerance. Each sweep is cheap and it scales to millions of sparse states, but convergence can be slow when the chain is likely to loop many times before absorbing.

import numpy as np

def step(s, roll):
    t = s + roll
    if t > 9: t = s            # exact roll needed
    if t == 7: t = 1           # chute
    return t

n = 9                          # states 0..8 are unknown, 9 is the goal
A, b = np.eye(n), np.ones(n)
for s in range(n):
    for roll in (1, 2, 3):
        t = step(s, roll)
        if t != 9:
            A[s][t] -= 1 / 3
E = np.linalg.solve(A, b)
print(E[0])                    # 9.18697...

v = [0.0] * 10                 # value iteration, same answer
while True:
    nv = [0.0] * 10
    for s in range(9):
        nv[s] = 1 + sum(v[step(s, r)] for r in (1, 2, 3)) / 3
    delta, v = max(abs(x - y) for x, y in zip(nv, v)), nv
    if delta < 1e-9: break
print(v[0])                    # 9.18697... after about 106 sweeps

Both methods give about 9.187 expected turns from square 0, and a 200,000-game simulation lands at 9.19, which is the kind of independent check you should run on any probability DP before trusting it. The system is solvable only if the goal is reachable from every state; if some state can never reach the goal, its expected time is infinite and the matrix is singular. Check reachability with a graph search first instead of discovering it as a numerical error.

Adding decisions: optimal stopping

Adding choices turns the recurrence into a Bellman equation: at each state, take the action whose expected value is best. A small but instructive example: you may roll a die up to three times and keep the last roll you choose to stop on. With one roll left, the value is 3.5. With two rolls left, keep a roll only if it beats 3.5, so you keep 4, 5 or 6 and reroll otherwise: (4+5+6)/6 + (3/6) * 3.5 = 4.25. With three rolls left, keep only 5 or 6: (5+6)/6 + (4/6) * 4.25 = 14/3, about 4.67.

from fractions import Fraction as F

def best_value(rolls_left):
    if rolls_left == 1:
        return F(7, 2)
    cont = best_value(rolls_left - 1)              # value of rolling again
    return sum(max(F(face), cont) for face in range(1, 7)) / 6

print(best_value(3))                               # 14/3

The policy falls out of the same table: stop when the current roll beats the continuation value. This is a reinforcement-learning value function with the model known exactly, and the same structure as expectimax, where max nodes alternate with chance nodes.

Numbers: fractions, modular inverses and log space

Probability DP lives or dies on number representation. There are three choices, and the right one depends on who consumes the answer.

Exact fractions (Python's Fraction) are ideal for a reference implementation and golden test values, but denominators grow fast, so keep them out of hot loops.

Modular arithmetic is how competitive programming and some cryptographic protocols report exact rationals: the answer p/q is printed as p * q^(MOD-2) mod MOD for a prime modulus, using Fermat's little theorem for the inverse. For the coupon collector, 147/10 modulo 998244353 is 99824450. Every division in the recurrence, including the 1 / (1 - k/6) from a self-loop, becomes a multiplication by a modular inverse, so precompute the inverses of the small denominators you use repeatedly. See modular exponentiation for the inverse computation.

Floating point is what real systems use, and it has two traps. Subtracting nearly equal quantities, as in 1 - p for p close to 1, loses digits; prefer computing the complement directly. Long products underflow: the probability of one particular sequence of 1,100 fair coin flips is 2^-1100, which is below the smallest positive double (about 2^-1074), so it rounds to exactly 0. Work in log space instead, where that probability is about -762.5, and add log-probabilities with the log-sum-exp trick so sums of tiny terms stay accurate.

import math

def logsumexp(xs):
    m = max(xs)
    if m == -math.inf:
        return m
    return m + math.log(sum(math.exp(x - m) for x in xs))

# forward step in log space: lp_next[t] = logsumexp(lp[s] + log p(s, t) over s)

From puzzles to production problems

Real problems reduce to the same pattern. Expected time to failure of a redundant system is an absorbing chain over how many parts are alive; a retry budget is a forward DP over attempts. Name the state, write the transition function, pick the direction, then pick the solver. With few states but many steps, the forward pass is a vector times the same matrix t times, so matrix exponentiation reaches 10¹⁸ steps in O(n³ log t). With enormous state spaces, switch to Monte Carlo and report a confidence interval.

Failure modes

• Non-Markov state. The transition reads a variable the state omits. Symptom: the DP disagrees with simulation by a stable margin. Fix: derive the state from the transition function.
• Wrong direction. Computing a forward distribution and then multiplying by values to get an expectation over a stopping time, without tracking time. Symptom: answers that are close but not right. Fix: use backward expectation DP for expected costs.
• Unsolved cycles. Memoised recursion over a cyclic graph recurses forever or, with a visited guard, returns a half-computed value. Fix: self-loop algebra, a linear solve or iteration.
• Probability leak. Overshoot or boundary handling drops mass. Fix: assert each forward layer sums to 1.
• Precision loss. Underflow in long products, cancellation in 1 - p. Fix: log space, direct complements and a fractional reference implementation for tests.

Trade-offs

What to do next

1. Re-derive the coupon collector table by hand for a four-sided die and check it against code (the answer is 25/3).
2. Take one problem you care about and write its transition function as a pure function before writing any DP; list every variable it reads as part of the state.
3. Add a layer-sum assertion to every forward DP and a simulation test to every backward DP.
4. Implement the chute board with both Gaussian elimination and value iteration, then add a second chute and watch the iteration count grow.
5. Port one recurrence to modular arithmetic and one to log space, so both representations are in your toolkit.
6. Continue with dynamic programming fundamentals and the animated DP introduction if state design still feels shaky.