Numpy数组广播优化多元正态分布二次型计算问询

针对你需要计算的单个样本 $x$ 与多组 $\mu_i$、$S_i$ 的二次型 $(x - \mu_i)^T S_i (x - \mu_i)$，以下是几种比迭代或取对角线更高效的实现方式：

通用情况（适用于任意正定矩阵 $S$） #

方法1：使用 np.einsum 直接指定求和逻辑 #

einsum 可以精准描述矩阵运算的维度收缩，代码简洁且效率高：

import numpy as np

# 给定数据
mu = np.array([[1,2,3], [4,5,6]])
S = np.array([np.eye(3)*3, np.eye(3)*5])
X = np.array([np.random.random(3*10)]).reshape(10, 3)

# 取单个样本x
x = X[0]
# 计算x与每个mu的差值，形状(2,3)
diff = x - mu
# 直接计算每个i对应的二次型，结果形状(2,)
result = np.einsum('ij,ijk,ik->i', diff, S, diff)

解释：'ij,ijk,ik->i' 表示对每个 $i$，计算 $\sum_j \sum_k diff[i,j] * S[i,j,k] * diff[i,k]$，正好对应二次型的展开式。

方法2：广播 + 逐元素相乘求和 #

利用矩阵乘法后的广播特性，简化运算：

diff = x - mu
# 先计算diff与S的矩阵乘法，得到(2,3)，再与diff逐元素相乘后按行求和
result = (diff @ S * diff).sum(axis=1)

解释：diff @ S 得到每个 $(x-\mu_i)S_i$ 的行向量，与 diff 逐元素相乘后，按行求和等价于行向量与 diff 的点积，即二次型结果。

特殊优化：当 $S$ 是对角矩阵时 #

你的示例中 $S$ 是对角矩阵（$\text{eye}(3)$ 乘系数），可以利用对角矩阵的特性进一步减少计算量：

diff = x - mu
# 提取每个S_i的对角线元素，形状(2,3)
diag_S = np.diagonal(S, axis1=1, axis2=2)
# 计算每个维度上的平方差乘以对应对角线元素，再求和
result = (diff**2 * diag_S).sum(axis=1)

这种方法避免了矩阵乘法，仅需元素级运算，速度比通用方法更快。

结果验证 #

可以对比迭代方法的结果，确认正确性：

# 迭代方法的结果
iter_result = np.array([(x - mu[i]) @ S[i] @ (x - mu[i]).T for i in range(len(mu))])
# 验证一致性
print(np.allclose(result, iter_result))  # 输出True

内容的提问来源于stack exchange，提问作者thesecond