EM算法

import random


class EM():
    def __init__(self):
        # 设状态值概率
        self.t = 0.6
        self.q = 0.5

    def E(self, i, j):
        pA = (self.t ** i) * ((1 - self.t) ** j)
        pB = (self.q ** i) * ((1 - self.q) ** j)

        d = pA / (pA + pB)
        b = pB / (pA + pB)
        # print('选择A的概率:{:.2f}，选择B的概率:{:.2f}'.format(d, b))

        cd = d * i
        ad = d * j
        # print('硬币A，正面朝上次数的期望值：{:.1f} 反面：{:.1f}'.format(cd, ad))

        cb = b * i
        ab = b * j
        # print('硬币B，正面朝上次数的期望值：{:.1f} 反面：{:.1f}'.format(cb, ab))

        return cd, ad, cb, ab

    def M(self):
        qq = 0
        tt = 0
        jj = 0
        cc = 0

        td = [[5, 5], [9, 1], [8, 2], [4, 6], [7, 3]] # 结果
        for i, j in td:
            a, b, c, d = self.E(i, j)
            qq += a
            tt += b
            jj += c
            cc += d

        print(qq, tt)
        print('A正面概率：{:.2f} 反面概率：{:.2f}'.format(qq/(qq+tt), tt/(qq+tt)))
        print(jj, cc)
        print('B正面概率：{:.2f} 反面概率：{:.2f}'.format(jj/(jj+cc), cc/(jj+cc)))
        self.t = round(qq/(qq+tt), 3)
        self.q = round(jj/(jj+cc), 3)
        print('-'*50)

    def run(self):
        self.s = None
        self.p = None
        while 1:
            if (self.s == self.t) and (self.p == self.q):
                break
            else:
                self.s = self.t
                self.p = self.q
                self.M()


EM().run()