1D Random Walk

1D_random_walk.py

"""
start at x, each step either walk +1 or -1. If we arrive at 0 or n, then we stop.
question: what's the probability to stop at 0 and n with starting position x?


theoretically: P[stop at 0] = (n-x)/n, P[stop at n] = x/n


simulation results in the format "x, (P[stop at 0], P[stop at n])":
2 (0.98032, 0.01968)
8 (0.92075, 0.07925)
16 (0.83919, 0.16081)
32 (0.6816, 0.3184)
64 (0.35897, 0.64103)
"""

import numpy as np

def test(x, n=100):
    if x == 0:
        return (0, -1)
    
    if x == n:
        return (-1, 0)
    
    s = np.array([x, x]).reshape([2,1])
    t = 0
    steps = 1000
    while True:
        res = s + np.cumsum(np.random.choice([1, -1], size=steps)).reshape([1,-1])
        
        end_at_zero = np.where(res[0] == 0)[0]
        end_at_n = np.where(res[1] == n)[0]
        
        if end_at_zero.size > 0 and end_at_n.size > 0:
            if end_at_zero[0] < end_at_n[0]:  # end at zero first
                return (t + 1 + end_at_zero[0], -1)
            else:
                return (-1, t + 1 + end_at_n[0])
        
        if end_at_zero.size > 0:
            return (t + 1 + end_at_zero[0], -1)
        
        if end_at_n.size > 0:
            return (-1, t + 1 + end_at_n[0])
        
        t += 1 + steps
        s = res[:, -1].reshape([2,1])
        assert s[0] > 0 and s[1] > 0
 
 
for x in [2, 8, 16, 32, 64]:
    trials = 100000
    total = list(map(lambda _: test(x), range(trials)))
    total = np.array(total)
    
    cnt_end_at_zero = np.sum(total[:, 0] > 0)
    cnt_end_at_n = np.sum(total[:, 1] > 0)
    res = (cnt_end_at_zero/trials, cnt_end_at_n/trials)
    print(x, res)