EM算法求解男女升高-Python

EM算法求解男女升高-Python

目前关于EM算法理论分析很全了，代码还比较少，一般就停留在最后的似然函数。

问题

现在一个班里有100个男生和女生。我们假定男生的身高服从正态分布，女生的身高则服从另一个正态分布。
求解两个正态分布的参数。

求解

根据最后得到的似然函数，对均值和方差求导，即可得到解析式。

与一般的正态分布最大似然估计相似，不同的只是加了一个权重。

代码

import numpy as np
from scipy.stats import norm

# 设置随机种子
np.random.seed(0)

# 真实参数
mu_m_true, sigma_m_true = 175, 7    # 男生身高的真实均值和标准差
mu_f_true, sigma_f_true = 165, 5    # 女生身高的真实均值和标准差

# 生成样本数据
male_heights = np.random.normal(mu_m_true, sigma_m_true, 50)
female_heights = np.random.normal(mu_f_true, sigma_f_true, 50)
heights = np.concatenate([male_heights, female_heights])

# 初始化参数
mu_m, sigma_m = 170, 10  # 男生身高的初始均值和标准差
mu_f, sigma_f = 160, 10  # 女生身高的初始均值和标准差

# EM算法
iterations = 100  # 设置迭代次数
for _ in range(iterations):
    # E步骤：计算每个样本属于男生或女生的概率
    prob_m = norm.pdf(heights, mu_m, sigma_m)
    prob_f = norm.pdf(heights, mu_f, sigma_f)
    weights_m = prob_m / (prob_m + prob_f)
    weights_f = prob_f / (prob_m + prob_f)

    # M步骤：更新参数
    mu_m = np.sum(weights_m * heights) / np.sum(weights_m)
    sigma_m = np.sqrt(np.sum(weights_m * (heights - mu_m)**2) / np.sum(weights_m))
    
    mu_f = np.sum(weights_f * heights) / np.sum(weights_f)
    sigma_f = np.sqrt(np.sum(weights_f * (heights - mu_f)**2) / np.sum(weights_f))

print(mu_m, sigma_m, mu_f, sigma_f)