条件概率与条件期望

A coin, having probability p of landing heads, is continually flipped until at least one head and one tail have been flipped.
(a) Find the expected number of flips needed.
(b) Find the expected number of flips that land on heads.
(c) Find the expected number of flips that land on tails.
(d) Repeat part (a) in the case where flipping is continued until a total of at least
two heads and one tail have been flipped.

Answer：
(a) $$\frac{p^2-p+1}{p(1-p)}$$
(b) $$\frac{p^2-p+1}{1-p}$$
(c) $$\frac{p^2-p+1}{p}$$
(d) $$\frac{1}{1-p}+\frac{(p+2)(1-p)}{p}$$

对d进行仿真：

import random
import matplotlib.pyplot as plt
import numpy as np

def sim(p):
    count = 0
    round = 100     #重复次数
    for i in range(round):
        head = 0
        tail = 0
        while True:
            #print('p = {}, head = {}, tail = {}'.format(p,head,tail))
            count = count+1
            if  random.random()<p:
                head = head+1
            else:
                tail = tail+1
            if head>=2 and tail>=1:   
                break
    return count/round

def fun(p):
    num = 1/(1-p) + (p+2)*(1-p)/p  #answer
    return num

x = np.linspace(0.01,0.99,20)
y_simulation = [sim(i) for i in x]
y_target = [fun(i) for i in x]
plt.plot(x,y_simulation,'b')   #蓝线代表仿真结果
plt.plot(x,y_target,'or')   #红点代表理论结果
plt.show()