找钥匙问题

小明把钥匙丢了，只知道钥匙肯定丢在家里或者学校里。如果丢在了家里，在家里找一次能够找到的概率为a；如果丢在了学校里，在学校里找一次能够找到的概率为b。最初小明认为钥匙丢在家里的概率为p。问：小明应该采取什么策略来找钥匙才能够最快找到钥匙（寻找次数尽量少）？期望寻找次数是多少？

a = 0.8  # 在家找到的概率
b = 0.4  # 在学校找到的概率


def solve(p, c):
    # 钥匙在家的概率为p，返回最优决策、期望找到钥匙所花费的步数
    if c == 0:
        return max(1 / a, 1 / b)  # 如果最后没有找到
    if p > b / (a + b):
        step = solve(p * (1 - a) / (1 - p * a), c - 1)
        ifhome = 1 + step * (1 - p * a)  # 如果去家里找，期望找到的步数
        return ifhome
    else:
        step = solve(p / (1 - (1 - p) * b), c - 1)
        ifschool = 1 + step * (1 - (1 - p) * b)
        # 选取两个儿子中期望步数较小者
        return ifschool


import pylab as plt
import numpy as np
print(solve(0.3, 200))
xs = np.linspace(0, 1, 1000)
ps = [solve(i, 200) for i in xs]
plt.plot(xs, ps)
plt.xlabel("p")
plt.ylabel("step")
print(b/(a+b))
plt.plot([b/(a+b)] * 100, np.linspace(0, max(ps), 100), c='r')
plt.show()