Probability and Expected Value in Competitive Programming

Sanjeev SharmaSanjeev Sharma
3 min read

Advertisement

Probability in Competitive Programming

Expected Value DP

Define E[state] = expected cost/steps to reach goal from state.

Classic: Expected dice rolls to reach 6

E = 1/6 * 1 + 5/6 * (1 + E)
E = 1 + 5E/6
E/6 = 1
E = 6

Python Implementation

from functools import lru_cache

# Expected steps for random walk on number line
# From position i, move left or right with equal prob
# Stop when you reach 0 or n

def expected_steps(n, p=0.5):
    # E[i] = 1 + p*E[i+1] + (1-p)*E[i-1], E[0]=E[n]=0
    # Solve linear system
    # For p=0.5: E[i] = i*(n-i) (symmetric random walk)
    return [i * (n - i) for i in range(n + 1)]

# Dice problem: expected number of rolls until all faces seen
def coupon_collector(n):
    # E[k] = n/(n-k) when k faces seen, need 1 more distinct
    # Total E = sum n/k for k = n, n-1, ..., 1 = n * H(n)
    import math
    H = sum(1/k for k in range(1, n+1))
    return n * H

print(coupon_collector(6))  # ~14.7 for a die

# Probability DP: survive game with N rounds, each round 1/6 chance to die
def survive_prob(rounds):
    p = 1.0
    for _ in range(rounds):
        p *= 5/6
    return p

# Geometric distribution: number of trials until first success
# P(X=k) = (1-p)^(k-1) * p
# E[X] = 1/p
def geometric_expected(p): return 1/p

C++ Implementation

#include <bits/stdc++.h>
using namespace std;

// DP with probability
// 837. New 21 Game
double new21Game(int n, int k, int maxPts) {
    if (k == 0 || n >= k + maxPts) return 1.0;
    vector<double> dp(n + 1, 0.0);
    dp[0] = 1.0;
    double wSum = 1.0; // window sum

    for (int i = 1; i <= n; i++) {
        dp[i] = wSum / maxPts;
        if (i < k) wSum += dp[i];
        if (i >= maxPts) wSum -= dp[i - maxPts];
    }

    double ans = 0;
    for (int i = k; i <= n; i++) ans += dp[i];
    return ans;
}

Java Implementation

public class ProbabilityDP {
    // LeetCode 688: Knight Probability
    public double knightProbability(int n, int k, int r, int c) {
        double[][] dp = new double[n][n];
        dp[r][c] = 1.0;
        int[][] dirs = {{-2,-1},{-2,1},{-1,-2},{-1,2},{1,-2},{1,2},{2,-1},{2,1}};

        for (int step = 0; step < k; step++) {
            double[][] next = new double[n][n];
            for (int i = 0; i < n; i++)
                for (int j = 0; j < n; j++) {
                    if (dp[i][j] == 0) continue;
                    for (int[] d : dirs) {
                        int ni = i + d[0], nj = j + d[1];
                        if (ni >= 0 && ni < n && nj >= 0 && nj < n)
                            next[ni][nj] += dp[i][j] / 8.0;
                    }
                }
            dp = next;
        }

        double res = 0;
        for (double[] row : dp) for (double v : row) res += v;
        return res;
    }
}

LeetCode Problems

  • 688. Knight Probability in Chessboard — probability DP
  • 837. New 21 Game — sliding window probability DP
  • 1227. Airplane Seat Assignment Probability — math insight (answer is always 0.5 for n>=2)
  • 808. Soup Servings — expected value with memoization

Advertisement

Sanjeev Sharma

Written by

Sanjeev Sharma

Full Stack Engineer · E-mopro