牛客题解 | 两个正态分布之间的KL散度

题目

题目链接

两个正态分布之间的KL散度（Kullback-Leibler Divergence）是一种衡量两个分布之间差异的指标，其计算公式为：

\[KL(p||q) = \int p(x) \log \frac{p(x)}{q(x)} dx \]

其中，\(p(x)\)和\(q(x)\)分别是两个正态分布的概率密度函数。
但是本题中，使用的计算公式是：

对于两个正态分布 \(p(x) = N(\mu_p, \sigma_p^2)\) 和 \(q(x) = N(\mu_q, \sigma_q^2)\)，它们之间的KL散度可以表示为：

\[KL(p||q) = \ln\frac{\sigma_q}{\sigma_p} + \frac{\sigma_p^2 + (\mu_p - \mu_q)^2}{2\sigma_q^2} - \frac{1}{2} \]

标准代码如下

def kl_divergence_normal(mu_p, sigma_p, mu_q, sigma_q):
    term1 = np.log(sigma_q / sigma_p)
    term2 = (sigma_p ** 2 + (mu_p - mu_q) ** 2) / (2 * sigma_q ** 2)
    kl_div = term1 + term2 - 0.5
    return kl_div