Probability and Expected-Value DP

This is an extension topic beyond most interviews; it pays off when a question asks for the probability of an outcome or the expected number of steps of a random process, which shows up in some hard OAs and quant-flavored rounds.

Intuition

A random process moves between states with known probabilities. Two kinds of questions come up, and they run in opposite directions:

• Probability DP (forward): P[state] = probability of being in that state after some steps. Start from the initial state with probability 1 and push probability mass along transitions: P[next] += P[cur] * p(cur -> next). Mutually exclusive ways of reaching a state add up; independent events along one way multiply.
• Expectation DP (backward): E[state] = expected remaining cost (steps, days, money) to reach the goal from that state. The goal has E = 0, and by linearity of expectation E[cur] = cost + sum of p(cur -> next) * E[next]. You compute it from the goal backward and the answer is E[start].

If a state can stay where it is with probability q, the equation contains E[cur] on both sides: E = 1 + q * E + rest, so E = (1 + rest) / (1 - q). Self-loops are fine; longer cycles need solving a linear system, which is beyond interview scope.

When to use it in interviews / OAs

• "Return the probability that ... after k moves / rounds" with a small state space (board cells, counts, scores).
• "Expected number of steps/days/draws until ..." where progress can be summarized by a few counters.
• Coins with different probabilities, random walks on a grid, games where both sides draw at random.

Template 1: forward probability on a grid

A knight starts on an n x n board and makes k moves, each uniformly random among its 8 jumps; once it leaves the board it stays off. What is the probability that it is still on the board?

def knight_probability(n: int, k: int, row: int, col: int) -> float:
    moves = [(1, 2), (2, 1), (2, -1), (1, -2), (-1, -2), (-2, -1), (-2, 1), (-1, 2)]
    # prob[r][c] = probability that the knight stands on (r, c) after the moves so far
    prob = [[0.0] * n for _ in range(n)]
    prob[row][col] = 1.0
    for _ in range(k):
        nxt = [[0.0] * n for _ in range(n)]
        for r in range(n):
            for c in range(n):
                if prob[r][c]:
                    for dr, dc in moves:
                        nr, nc = r + dr, c + dc
                        if 0 <= nr < n and 0 <= nc < n:
                            nxt[nr][nc] += prob[r][c] / 8
        prob = nxt
    return sum(map(sum, prob))

Mass that jumps off the board is simply dropped, so the remaining total is the answer.

Template 2: counting with probabilities

Coins have different head probabilities p[i]. What is the probability of exactly target heads? It is a 0/1 knapsack where "take" has weight p and "skip" has weight 1 - p:

def prob_exact_heads(probs: list[float], target: int) -> float:
    # f[j] = probability of exactly j heads among the coins tossed so far
    f = [1.0] + [0.0] * target
    for p in probs:
        for j in range(target, -1, -1):  # descending, exactly like a 0/1 knapsack
            f[j] = f[j] * (1 - p) + (f[j - 1] * p if j else 0.0)
    return f[target]

Template 3: backward expectation with a self-loop

Every day you find one bug. It belongs to one of n categories and one of s subsystems, each uniformly at random and independently. What is the expected number of days until every category and every subsystem has at least one bug? The progress is summarized by (i, j): categories seen and subsystems seen.

From (i, j), tomorrow's bug is old in both coordinates with probability (i/n) * (j/s), which keeps you in the same state. The other three outcomes advance i, j or both.

def expected_days(n: int, s: int) -> float:
    # E[i][j] = expected extra days when i categories and j subsystems are already covered
    E = [[0.0] * (s + 1) for _ in range(n + 1)]
    for i in range(n, -1, -1):
        for j in range(s, -1, -1):
            if i == n and j == s:
                continue  # goal reached: no more days needed
            stay = (i / n) * (j / s)  # nothing new today: same state tomorrow
            rest = 1.0  # today's day
            if i < n:
                rest += (1 - i / n) * (j / s) * E[i + 1][j]
            if j < s:
                rest += (i / n) * (1 - j / s) * E[i][j + 1]
            if i < n and j < s:
                rest += (1 - i / n) * (1 - j / s) * E[i + 1][j + 1]
            E[i][j] = rest / (1 - stay)  # solve E = rest + stay * E
    return E[0][0]

With s = 1 this is the coupon collector problem, and the result matches the closed form n * (1 + 1/2 + ... + 1/n): a quick way to sanity-check the recurrence.

Complexity

States times transitions, as usual: knight O(k * n^2 * 8), coins O(len(probs) * target), bugs O(n * s). Memory can be reduced to one layer for the forward versions.

Common pitfalls

• Wrong direction: probabilities flow forward from the start; expectations are easiest backward from the goal. Mixing them up gives equations that reference states not yet computed.
• Forgetting the self-loop: writing E = 1 + q * E + ... and then using the old value of E on the right side. Solve for E algebraically first.
• Adding instead of multiplying: independent events along one path multiply; different disjoint paths into the same state add.
• Probabilities that do not sum to 1: check that the outgoing probabilities of each state add up to 1 (minus whatever deliberately leaves, like jumping off the board).
• Floating-point comparisons: answers are usually accepted within 1e-5; never test floats with ==. Use fractions.Fraction to verify small cases exactly.
• Huge parameters: some problems converge (the probability tends to 1 for very large inputs); look for that cutoff instead of building an enormous table.

Interview variations

• Knight or robot staying on a board: forward probability over cells.
• Draw cards until the total reaches a threshold (the "21 game" pattern): forward probability where P[x] is a sum over a window of previous values, kept as a sliding-window sum.
• Serving two soups with random portions: memoized probability over the two remaining amounts, with a cutoff for large inputs.
• Exactly k heads with biased coins: probability knapsack.
• Expected steps of a random walk or expected rolls until all faces appear: backward expectation, often with self-loops.
• Position of a frog jumping randomly on a tree after t seconds: forward probability along tree edges.