PaC Exercise 3.6

import numpy as np

Exercise 3.6

For a coin that comes up head independently with probability p on each flip, what is the variance in the number of flips until the kth head appears?

Solution: Let X be the number of flips until the kth head appears and let Y_i, i = 1, …, k, be the number of throws until we see the ith head. Then each Y_i is geometric distributed with parameter p and all Y_is are independent.

Therefore,

$$ \mathrm{Var}(X) = \sum_{i=1}^k \mathrm{Y_i} = k \frac{(1-p)}{p^2}, $$

For example, for k = 11 and p = 0.35 this yields 21.224…. We verify this by simple sampling.

num_samples = 5*10**6
k= 4
p = 0.35

samples = np.random.choice(a=[0,1], p=[1-p, p], size=num_samples)

ith_heads = np.cumsum(samples) % k

seq = np.array([k-1, 0])

indices = np.where([np.array_equal(a, seq) 
                    for a in np.stack([ith_heads[:-1], ith_heads[1:]], axis=1)])  # where do we see the $k$th head?
np.diff(np.c_[0, indices]).var()

21.133861223328264