【机器学习系列】EM算法求解三硬币问题（python版本）

三硬币模型

假设有3枚硬币,分别记作A,B,C。这些硬币正面出现的概率分别为π，p和q。投币实验如下，先投A，如果A是正面，即A=1，那么选择投B；A=0，投C。最后，如果B或者C是正面，那么y=1；是反面，那么y=0；独立重复n次试验（n=10)，观测结果如下： 1,1,0,1,0,0,1,0,1,1假设只能观测到投掷硬币的结果，不能观测投掷硬币的过程。问如何估计三硬币正面出现的概率，即π，p和q的值。

参考李航《统计学习方法》P155“三硬币模型”

推导过程

参考“三硬币问题”http://www.crescentmoon.info/?p=573

代码（运行环境python2.7.3）

#coding:gbk

import math

[pi,p,q]=[0.4,0.6,0.7]

x=[1,1,0,1,0,0,1,0,1,1]

def cal_u(pi1,p1,q1,xi):
    return pi1*math.pow(p1,xi)*math.pow(1-p1,1-xi)/float(pi1*math.pow(p1,xi)*math.pow(1-p1,1-xi)+(1-pi1)*math.pow(q1,xi)*math.pow(1-q1,1-xi))

def e_step(pi1,p1,q1,x):
    return [cal_u(pi1,p1,q1,xi) for xi in x]

def m_step(u,x):
    pi1=sum(u)/len(u)
    p1=sum([u[i]*x[i] for i in range(len(u))]) / sum(u)
    q1=sum([(1-u[i])*x[i] for i in range(len(u))]) / sum([1-u[i] for i in range(len(u))])
    return [pi1,p1,q1]

def run(start_x,start_pi,start_p,start_q,iter_num):
    for i in range(iter_num):
        u=e_step(start_pi,start_p,start_q,x)
        print i,[start_pi,start_p,start_q]
        if [start_pi,start_p,start_q]==m_step(u,x):
            break
        else:
            [start_pi,start_p,start_q]=m_step(u,x)

run(x,pi,p,q,100)