Probability and Expected Value DP: From Coin Flips to Knight Walks

Algorithm/Topic Statement#

Probability and expected value problems ask you to compute the long-run average or the chance of an event in a system that evolves randomly. In a coding interview these problems show up as expected number of dice rolls, probability of a knight surviving k moves on a chessboard, the chance of a card game ending with a specific hand, or the expected time for a random walk to reach a barrier. The technique that solves nearly all of them is probability DP, which sets up a state, writes a recurrence weighting transitions by their probabilities, and applies memoization or tabulation. The math grounding comes from elementary probability: linearity of expectation, conditional probability, and the law of total expectation.

Why This Topic Matters#

Quantitative finance firms like Jane Street, Two Sigma, and Citadel build entire interview rounds around these problems. The classic dice and coin questions are warmups for systems thinking that traders use every day. Even at general tech companies, expected value DP appears whenever a problem involves randomness, like the New 21 Game on LeetCode or any algorithm using reservoir sampling. Beyond interviews, the same techniques drive randomized algorithms, queueing theory, and reinforcement learning. Knowing how to translate a verbal probability puzzle into a recurrence is a skill that pays dividends in research and engineering. It also forces clean reasoning about state, conditioning, and base cases, which transfers to nearly every other DP topic.

The Core Insight (math intuition + proof sketch)#

The key tool is linearity of expectation: the expected value of a sum equals the sum of expected values, even when the underlying random variables are dependent. This means you can decompose a complicated random process into pieces, compute the expectation of each, and add them together. The second tool is the law of total expectation, which says the expectation of X equals the weighted sum, over all possible immediate outcomes, of the conditional expectation given that outcome. Translated to code, this becomes a one-line recurrence: expected steps from state S equals one plus the weighted average of expected steps from every reachable next state, where the weights are transition probabilities.

The proof sketch for solving expected dice rolls to see a six is illustrative. Let E be the expected number of rolls. With probability one sixth you stop on this roll, contributing one. With probability five sixths you roll once and then face the same problem again, contributing one plus E. So E equals one sixth times one plus five sixths times the quantity one plus E. Expanding gives E equals one plus five sixths E. Subtract to get one sixth E equals one, hence E equals six. This recurrence pattern, namely E equals one plus weighted future Es, generalizes to almost every Markov chain expectation.

For probability, the recurrence is similar but starts from one instead of zero and propagates forward. Probability of being in state S at time t equals the sum, over predecessors P, of the probability of being in P at time t minus one times the transition probability from P to S. This is a finite-horizon Markov chain.

Visual Dry Run / Worked Example#

Take the LeetCode knight probability problem: a knight starts at position r, c on an n by n board and makes k moves uniformly at random. Compute the probability it remains on the board after all k moves. Build a 2D probability grid, dp, with dp at the start position equal to one and all others zero. For each step, create a new grid called next. For every cell with positive probability, distribute one eighth of that probability to each of the eight knight moves that land inside the board.

For n equal to 3, r equal to zero, c equal to zero, and k equal to two, the first step distributes the starting probability one to two valid landing squares, each receiving one eighth. The remaining six moves leave the board and contribute nothing. So after one step the total probability on the board is two eighths or one quarter. The second step takes each of those two squares and distributes one eighth of their probability to in-board squares. Some squares receive contributions from both predecessors. Sum over all in-board cells at the end to get the answer, which is six over sixty-four or three over thirty-two.

Verify the math by enumerating eight times eight equals sixty-four equally likely two-move sequences and counting how many keep the knight on the board.

Solution / Implementation#

Python (probability DP and expected value DP)#

from functools import lru_cache
 
def knight_probability(n, k, r, c):
    moves = [(-2,-1),(-2,1),(-1,-2),(-1,2),(1,-2),(1,2),(2,-1),(2,1)]
    dp = [[0.0]*n for _ in range(n)]
    dp[r][c] = 1.0
    for _ in range(k):
        nxt = [[0.0]*n for _ in range(n)]
        for i in range(n):
            for j in range(n):
                if dp[i][j] == 0:
                    continue
                for dx, dy in moves:
                    ni, nj = i+dx, j+dy
                    if 0 <= ni < n and 0 <= nj < n:
                        nxt[ni][nj] += dp[i][j] / 8.0
        dp = nxt
    return sum(sum(row) for row in dp)
 
def expected_dice_rolls_until_six():
    return 6.0
 
def coupon_collector(n):
    return n * sum(1.0 / k for k in range(1, n + 1))

JavaScript#

function knightProbability(n, k, r, c) {
  const moves = [[-2,-1],[-2,1],[-1,-2],[-1,2],[1,-2],[1,2],[2,-1],[2,1]];
  let dp = Array.from({length: n}, () => new Array(n).fill(0));
  dp[r][c] = 1;
  for (let s = 0; s < k; s++) {
    const next = Array.from({length: n}, () => new Array(n).fill(0));
    for (let i = 0; i < n; i++) {
      for (let j = 0; j < n; j++) {
        if (dp[i][j] === 0) continue;
        for (const [dx, dy] of moves) {
          const ni = i + dx, nj = j + dy;
          if (ni >= 0 && ni < n && nj >= 0 && nj < n) {
            next[ni][nj] += dp[i][j] / 8;
          }
        }
      }
    }
    dp = next;
  }
  let total = 0;
  for (const row of dp) for (const v of row) total += v;
  return total;
}

Time complexity for the knight problem is order k times n squared times eight. Space is order n squared per layer. Memoization with hashed states reduces redundant work for problems whose state space is sparse.

Common Mistakes#

A frequent error is double counting probabilities by adding from a stale grid while updating in place. Always allocate a fresh next grid each step. Another trap is forgetting that probabilities must remain nonnegative and may sum to less than one when the knight can leave the board, which is precisely the answer you want. Many candidates write the recurrence backward, conditioning on the previous state instead of the next, and then mix up the direction of summation. For expected value problems, the most common mistake is omitting the plus one for the cost of the current step. Finally, watch for self-loops in the state graph. If a transition can return to the same state with positive probability, the equation has the unknown on both sides and you must solve algebraically rather than relying on memoization, which would loop forever.

Interview Tips#

When the problem says expected, picture a recurrence that opens with one plus a weighted average. When the problem says probability, picture forward propagation of a distribution. State the random variable explicitly before writing code, for example let dp i, j be the probability of being at i, j after some step. If memoization is natural, use it; if the state space is dense, use tabulation. Mention linearity of expectation aloud when you decompose an expectation into pieces. If the interviewer presses on convergence or self-loops, derive the algebra by hand. These small narrations communicate that you understand the math, not just the recipe.

Follow-up Questions#

How would you compute the expected number of distinct values seen after k draws with replacement from a uniform distribution? Could you derive the variance, not just the expectation, in a probability DP? What changes when transitions have nonuniform probabilities, like a biased coin? Can you implement an exact rational arithmetic version using Python fractions to avoid floating point error in long DPs?

Key Takeaways#

• Probability DP defines a state and writes a recurrence whose terms are weighted by transition probabilities.
• Expected value recurrences typically open with one plus a weighted average over reachable next states.
• Linearity of expectation lets you decompose complex random variables into independent or correlated pieces with the same sum rule.
• Always allocate a fresh next grid when propagating probabilities to avoid double counting.
• Self-loops require algebraic resolution rather than blind memoization.
• These techniques are required by quantitative finance interviews and apply to randomized algorithms across the field.