考虑简单的抛硬币问题. 现有两枚硬币$A$和$B$，正面朝上的概率分别为$\theta_A, \theta_B$, 结果朝上 记为H~(head), 朝下记为T~(tail). 独立地进行$N$轮实验, 在第$k$轮实验中, 以均等概率选择一枚硬币$Z_k \in \{A,B\}$并重复抛掷$M$次, 其中硬币朝上的次数$X_k$为可观测变量, 而选择的硬币类型$Z_k$为隐变量不可观测. 我们将使用EM算法, 迭代一次, 对参数$\theta = (\theta_A, \theta_B)$进行估计, 使用的实验数据如表\ref{table: EM}所示. 具体而言共3轮实验, 每轮选取的硬币记为$z_i ~ (i=1,2,3)$, 抛掷10次并记录结果, 硬币朝上的次数记为$x_i ~ (i=1,2,3)$. 假设参数的初始值$\theta^0 = (0.6, 0.5)$. 请结合实验数据, 推断出隐变量取值$\z = (z_1, z_2)$的分布, 即推断出第$i$轮实验~($i=1,2,3$)中抛掷硬币$A$、硬币$B$各自的概率

首先，我们需要构建完整数据的似然函数，即将隐变量$\z$也考虑进去，得到完整数据似然函数为：

$$
L_c(\theta) = \prod_{i=1}^N p(\x_i, \z_i | \theta) = \prod_{i=1}^N [p(\x_i | \z_i, \theta) \cdot p(\z_i)]
$$

其中，$p(\x_i | \z_i, \theta)$是已知选择了硬币$z_i$的情况下，观测到$x_i$的概率，即：

$$
p(\x_i | \z_i, \theta) = \binom{M}{x_i} \theta_{z_i}^{x_i}(1-\theta_{z_i})^{M-x_i}
$$

$p(\z_i)$是隐变量$\z_i$的先验分布，由于每次实验选择硬币的概率均等，因此：

$$
p(\z_i) = \begin{cases}
0.5 & \text{if } \z_i = A \\
0.5 & \text{if } \z_i = B \\
\end{cases}
$$

对于EM算法，我们需要先求出完整数据的对数似然函数，即：

$$
\begin{aligned}
\ln L_c(\theta) &= \sum_{i=1}^N [\ln p(\x_i | \z_i, \theta) + \ln p(\z_i)] \\
&= \sum_{i=1}^N \left[\ln \binom{M}{x_i} + x_i \ln \theta_{z_i} + (M-x_i)\ln(1-\theta_{z_i}) + \ln 0.5 \right]
\end{aligned}
$$

接下来，我们需要求在给定参数$\theta$下，隐变量$\z$的后验概率$p(\z_i | \x_i, \theta)$。根据贝叶斯公式：

$$
p(\z_i | \x_i, \theta) = \frac{p(\x_i, \z_i | \theta)}{p(\x_i | \theta)} = \frac{p(\x_i | \z_i, \theta) \cdot p(\z_i)}{\sum_{\z_i} p(\x_i | \z_i, \theta) \cdot p(\z_i)}
$$

其中，分母是边缘概率，可以按照以下方式计算：

$$
p(\x_i | \theta) = \sum_{\z_i} p(\x_i | \z_i, \theta) \cdot p(\z_i)
$$

然后，我们可以根据EM算法的求解步骤，依次进行E步和M步：

1. E步：在给定参数$\theta^t$的情况下，计算隐变量$\z$的后验概率$p(\z_i | \x_i, \theta^t)$。
2. M步：在给定隐变量的后验概率$p(\z_i | \x_i, \theta^t)$的情况下，计算参数的更新值$\theta^{t+1}$。

具体而言，在E步中，我们需要计算每个隐变量的后验概率：

$$
\begin{aligned}
p(z_i=A | \x_i, \theta^t) &= \frac{p(\x_i | z_i=A, \theta^t) \cdot p(z_i=A)}{p(\x_i | \theta^t)} \\
&= \frac{\binom{M}{x_i} \theta^t_A}{\binom{M}{x_i} \theta^t_A + \binom{M}{x_i} \theta^t_B} \\
&= \frac{\theta^t_A}{\theta^t_A + \theta^t_B}
\end{aligned}
$$

$$
\begin{aligned}
p(z_i=B | \x_i, \theta^t) &= \frac{p(\x_i | z_i=B, \theta^t) \cdot p(z_i=B)}{p(\x_i | \theta^t)} \\
&= \frac{\binom{M}{x_i} \theta^t_B}{\binom{M}{x_i} \theta^t_A + \binom{M}{x_i} \theta^t_B} \\
&= \frac{\theta^t_B}{\theta^t_A + \theta^t_B}
\end{aligned}
$$

在M步中，我们需要计算参数的更新值。对于硬币$A$，我们可以将完整数据的对数似然函数对$\theta_A$求导数，并令其等于$0$，得到：

$$
\frac{\partial \ln L_c(\theta)}{\partial \theta_A} = \sum_{i=1}^N\left[\frac{x_i}{\theta_A} - \frac{M-x_i}{1-\theta_A}\right] = 0
$$

化简后得到：

$$
\theta_A^{t+1} = \frac{\sum_{i=1}^N p(z_i=A | \x_i, \theta^t) x_i}{\sum_{i=1}^N p(z_i=A | \x_i, \theta^t) M}
$$

同理，对于硬币$B$，我们可以得到：

$$
\theta_B^{t+1} = \frac{\sum_{i=1}^N p(z_i=B | \x_i, \theta^t) x_i}{\sum_{i=1}^N p(z_i=B | \x_i, \theta^t) M}
$$

根据以上推导，我们可以使用EM算法对参数$\theta$进行估计，具体实现过程可以参考下面的Python代码：