如何用Scipy稀疏矩阵与3D NumPy数组高效实现指定求和公式？

我有如下矩阵：

import scipy.sparse as sp
import numpy as np

a = sp.random(150, 150)
x = np.random.normal(0, 1, size=(150, 20))

需要实现以下公式：
$$\sigma_{k} = \sum_{ij}^{N} A_{ij} (x_{ik} - x_{jk})^2$$

我原本尝试计算差值：

diff = (x[:, None, :] - x[None, :, :]) ** 2
diff.shape  # -> (150, 150, 20)

a.shape  # -> (150, 150)

如果把稀疏矩阵a转为密集矩阵，可以直接用 einsum 计算：

np.einsum("ij,ijk->k", a.toarray(), (x[:, None, :] - x[None, :, :]) ** 2)

但a是规模极大的稀疏矩阵，转密集矩阵不可行；同时生成(150,150,20)的diff数组也会导致内存问题。不想用循环遍历，有没有基于NumPy的高效实现方式？

核心思路：公式拆解规避内存瓶颈 #

直接计算差值平方会生成三维巨型数组，稀疏矩阵转密集也会占用海量内存。我们可以通过展开平方项，将原公式转化为稀疏矩阵友好的运算形式：

展开平方项：
$$(x_{ik} - x_{jk})^2 = x_{ik}^2 - 2x_{ik}x_{jk} + x_{jk}^2$$

代入原求和式后拆分为三个独立项：
$$\sigma_k = \sum_{ij}A_{ij}x_{ik}^2 - 2\sum_{ij}A_{ij}x_{ik}x_{jk} + \sum_{ij}A_{ij}x_{jk}^2$$

代码实现 #

1. 计算稀疏矩阵的行和与列和：

row_sums = a.sum(axis=1).A.flatten()  # 行和，形状(150,)
col_sums = a.sum(axis=0).A.flatten()  # 列和，形状(150,)

2. 计算第一项与第三项：

x_squared = x ** 2
term1 = row_sums @ x_squared  # 结果形状(20,)
term3 = col_sums @ x_squared  # 结果形状(20,)

3. 计算第二项：
   $\sum_{ij}A_{ij}x_{ik}x_{jk}$ 等价于向量$x[:,k]$与稀疏矩阵a的二次型，即$x[:,k]^T @ a @ x[:,k]$，对所有k取对角线值即可：

xTaTx = x.T @ a @ x  # 稀疏矩阵乘法，结果形状(20,20)
term2 = -2 * np.diag(xTaTx)  # 取对角线元素，形状(20,)

4. 合并得到最终结果：

sigma = term1 + term2 + term3

优势总结 #

• 全程无需转换稀疏矩阵为密集矩阵，充分利用Scipy稀疏矩阵的高效运算特性
• 无巨型中间数组生成，内存占用仅与x和稀疏矩阵的非零元素数量相关
• 所有操作均为向量/矩阵级别的原生运算，比循环遍历效率提升显著

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