cswork / libraryOpen practice workspace

Algorithm tutorials

Probability and Expected Value for Coding Interviews

The probability toolkit interviews actually use: conditional probability and Bayes' rule, linearity of expectation with indicator variables, geometric waiting times and the coupon collector, first-step analysis for expected times, rejection sampling, and exact or simulated sanity checks.

Probability questions in coding interviews rarely need heavy theory. A handful of tools solve almost all of them: conditional probability and Bayes' rule, linearity of expectation with indicator variables, the geometric "wait until success" distribution, first-step analysis for expected times, and a quick simulation or exact enumeration to check the answer.

Intuition

  • A sample space lists the possible outcomes; an event is a set of outcomes; with equally likely outcomes, P(A) = |A| / |sample space|.
  • Conditional probability: P(A | B) = P(A and B) / P(B), the probability of A once you know B happened. Independence means P(A and B) = P(A) * P(B).
  • Total probability: split by cases, P(A) = sum of P(A | Bi) * P(Bi) over disjoint cases Bi. Bayes' rule flips a conditional: P(B | A) = P(A | B) * P(B) / P(A).
  • A random variable assigns a number to each outcome; its expectation is the probability-weighted average.
  • Linearity of expectation: E[X + Y] = E[X] + E[Y] always, even when X and Y are dependent. This is the single most useful fact.
  • Indicator variables: I_A is 1 if A happens and 0 otherwise, and E[I_A] = P(A). Write a count as a sum of indicators and its expectation is a sum of probabilities.
  • Geometric waiting: if each trial succeeds independently with probability p, the expected number of trials until the first success is 1 / p.
  • First-step analysis (total expectation): condition on what happens in the first step and write an equation for the expected remaining time.

When to use it in interviews / OAs

  • "What is the expected number of ...?" Try linearity with indicators before anything else.
  • "How many tries on average until ...?" Geometric waiting times, or first-step equations.
  • "Given that the test is positive, what is the chance ...?" Bayes' rule with total probability in the denominator.
  • Building a uniform generator from another one (rejection sampling), or analyzing randomized algorithms.
  • Computing exact answers as fractions, or modulo a prime as P * Q^-1 (see modular arithmetic).

Template: linearity and indicators

The expected number of fixed points of a random permutation is 1 for every n: each position is fixed with probability 1/n, and there are n positions. Brute force with exact fractions is the best way to convince yourself (and to test formulas):

from fractions import Fraction
from itertools import permutations


def expected_fixed_points(n: int) -> Fraction:
    # exact average over all n! equally likely permutations (small n only)
    perms = list(permutations(range(n)))
    total = sum(sum(p[i] == i for i in range(n)) for p in perms)
    return Fraction(total, len(perms))


def expected_distinct(n_bins: int, m_balls: int) -> Fraction:
    # balls thrown uniformly into bins; one indicator per bin:
    # P(bin is used) = 1 - P(every ball misses it) = 1 - (1 - 1/n)^m
    return n_bins * (1 - (1 - Fraction(1, n_bins)) ** m_balls)


def coupon_collector(n: int) -> Fraction:
    # with i kinds already seen, a new kind appears with probability (n - i) / n,
    # so the wait for it is geometric with mean n / (n - i)
    return sum(Fraction(n, n - i) for i in range(n))

The coupon collector's expectation is n * (1 + 1/2 + ... + 1/n), roughly n ln n: collecting all 50 kinds takes about 225 draws on average.

Template: first-step analysis

Expected flips until k heads in a row, with heads probability p. Let E[j] be the expected flips to first reach a run of j. To get from a run of j - 1 to a run of j you flip once more: with probability p you are done, otherwise you restart from zero and need E[j] more flips. So E[j] = E[j-1] + 1 + (1 - p) * E[j], i.e. E[j] = (E[j-1] + 1) / p.

def expected_flips_for_run(k: int, p: Fraction = Fraction(1, 2)) -> Fraction:
    e = Fraction(0)  # E[0] = 0: a run of length 0 is already there
    for _ in range(k):
        e = (e + 1) / p
    return e

For a fair coin that is 2^(k+1) - 2: 2 flips for one head, 6 for two in a row, 14 for three.

Template: Bayes' rule

A condition affects 1% of people; the test detects it 99% of the time and gives a false positive 5% of the time. Given a positive test, the chance of having the condition is only about 17%, because healthy people vastly outnumber sick ones:

def posterior(prior: float, sensitivity: float, false_positive: float) -> float:
    # P(condition | positive) by Bayes' rule
    p_positive = prior * sensitivity + (1 - prior) * false_positive  # total probability
    return prior * sensitivity / p_positive

Template: rejection sampling

Build a uniform 1..10 generator from a uniform 1..7 generator. Two rolls give 49 equally likely outcomes; keep the first 40 (four copies of 1..10) and retry otherwise. Every kept outcome is equally likely, so the result is uniform.

def rand10(rand7) -> int:
    while True:
        idx = (rand7() - 1) * 7 + rand7()  # uniform in 1..49
        if idx <= 40:
            return (idx - 1) % 10 + 1

Each attempt succeeds with probability 40/49, so the expected number of attempts is 49/40 and the expected number of rand7 calls is 2 * 49/40 = 2.45: a geometric waiting time again.

Complexity

The formulas are O(n) or O(k) arithmetic operations. Exact enumeration is exponential and only for checking small cases; Monte Carlo simulation with T trials has error around 1 / sqrt(T), so 10^4 to 10^5 trials give two or three correct digits, enough to catch a wrong formula.

Common pitfalls

  • Assuming independence: E[XY] = E[X] E[Y] needs independence; E[X + Y] = E[X] + E[Y] does not. Do not mix them up.
  • Non-equally-likely outcomes: counting outcomes only works when they are equally likely ("two boys" problems, unfair dice).
  • Conditioning on the wrong event: write down exactly what is known before computing P(A | B).
  • Forgetting the restart: in first-step analysis the failure branch often returns to an earlier state, so the unknown appears on both sides of the equation.
  • Floating-point answers: use Fraction for exact checks; compare floats with a tolerance.
  • Expected value vs typical value: the expectation may never occur in a single run (expected dice value 3.5).

Interview variations

  • Expected number of fixed points / cycles / records in a random permutation: indicators (records: 1 + 1/2 + ... + 1/n).
  • Coupon collector and expected draws to see every face of a die.
  • Expected flips until a pattern (HH vs HT have different answers: 6 and 4).
  • Implement rand10 using rand7, shuffle an array uniformly, pick a random element weighted by value: uniformity arguments.
  • Bayes questions about tests, spam filters or defective machines.
  • Expected value reported modulo 10^9 + 7: compute the fraction and multiply by the inverse of the denominator.