Mathematical Expectation for Interviews

In this post, I will just solve for expected value of a probabilistic model with as many methods as I can. You can encountered these types of problem in data science and quant interviews.

Problem

Assume there is a 2×2 grid, as shown in the figure below. You can randomly walk to a neighboring block (non-diagonal) with a 0.5 probability. For example, if you are at state 0, you can randomly move to state 1 with a probability 0.5 and state 2 with a probability 0.5. Question is, if I randomly start walking from state 0, what is the expected number of steps in which I will reach red state (state 3) for the first time?

Solution 1: Mean Equations

Let $x$ be expected number of steps from state 0 to reach state 3 for first time and $y$ be expected number of steps to reach from state 1. The expected number of steps from state 2 will also be $y$ as state 1 and state 2 are symmetric for the problem. (You can try by not making them same)

Now, if you start from step 0, there is 0.5 chance you will goto state 1 (in one step) and from there it will take you $y$ steps to reach state 3. Similarly, there is 0.5 chance you will goto state 2 and from there it will take $y$ time to reach state 3.

$$ x = 0.5 \times (1+y) + 0.5 \times (1+y) = 1 + y $$

Similarly, from state 1, there is 0.5 probability you will reach state 3 (in one step) and 0.5 chance you will reach state 0, from where it will again take $x$ steps to reach state 3.

$$ y = 0.5 \times (1) + 0.5 \times (1+x) =1 + 0.5x $$

Solving these equations you will get, $x=4, y=3$. The expected number of steps to reach state 3 starting from state 0 is 4

Solution 2: Expectation Calculation

In this method, we will simply expand expectation definition, and try to solve with it.

$$ E[X] = \sum_{X} X P(X)$$

If you start from state 0, we will always reach at state 3 in even steps. Try it.

Now,

$P(0)=0$, because if you never move, you will never reach there.

$P(2) = 2 \times 0.5^2$, because there is two ways to reach step 3 from step 0 in 2 steps, both with probability $0.5^2$: 0 → 1 → 3 and 0 → 2 → 3

$P(4) = 4 \times 0.5^4$

$P(6) = 8 \times 0.5^6$

$E[X] = 0 \times 0 + 2 \times 0.5 + 4 \times 0.5^2 + 6 \times 0.5^3 + …  $

The above expression is an AGP (Arthemetic and Geometric Progression) with $a = 1, d = 1$ and $r = 0.5$. The above sum resolves to 4.

Solution 3: Walk Simulation

Coding time. Let’s say you can simulate this walk. In such simulation you can get number of steps the agent took to reach state 3 for first time. The answer to our question is the average of such simulations.

from random import random as rand
def number_of_steps(current_state=0):
    state = [0,1,2,3]  # 3 is terminal, 0 is diagonal
    time_taken = 0
    while current_state != 3:
        if current_state == 0:
            current_state = 1 if rand() < 0.5 else 2
        elif current_state == 1:
            current_state = 0 if rand() < 0.5 else 3
        elif current_state == 2:
            current_state = 0 if rand() < 0.5 else 3
        time_taken += 1
    return time_taken

monte_carlo = [number_of_steps(0) for _ in range(1_000_000)]
result = sum(monte_carlo)/len(monte_carlo)

Run this in Python, you will find the result is close to 4.

Closing

Hope you found this interesting. Anyone your favorite? Let me know in the comments below. Is there any other method I have missed? I think we can formulate this problem as markov chain and calculate mean time spent in transient state.