Python中如何无循环实现稀疏加权向量点积矩阵计算？

核心解法 #

不用循环的关键是结合SciPy稀疏矩阵特性和NumPy向量化运算，两种方案按需选择：

方案1：基于Gram矩阵的稀疏逐元素乘法

先计算X的Gram矩阵（所有行向量的点积结果矩阵），这一步是底层优化的向量化操作，速度远超Python循环：

G = X @ X.T

如果S是SciPy的稀疏矩阵（比如csr_matrix/coo_matrix），直接用稀疏矩阵的逐元素乘法，它会自动只处理S的非零位置，完全不用手动遍历：

A = S.multiply(G)

方案2：内存友好版（避免生成全量Gram矩阵）

如果N特别大，全量Gram矩阵G会占用O(N²)内存，容易出现内存不足的情况。这时直接利用S的非零坐标批量计算：

1. 提取S的非零元素的行、列索引：

   rows, cols = S.nonzero()
   

2. 批量提取对应的X行向量，计算点积：

   dot_products = (X[rows] * X[cols]).sum(axis=1)
   

3. 用这些结果构建新的稀疏矩阵A：

   from scipy import sparse
   A = sparse.coo_matrix((dot_products, (rows, cols)), shape=S.shape)
   

代码示例 #

import numpy as np
from scipy import sparse

# 模拟测试数据
N, d = 1000, 100
X = np.random.randn(N, d)
# 生成稀疏矩阵S（仅1%非零元素）
S = sparse.random(N, N, density=0.01, format='coo')

# 方案1实现
G = X @ X.T
A1 = S.multiply(G)

# 方案2实现
rows, cols = S.nonzero()
dot_vals = (X[rows] * X[cols]).sum(axis=1)
A2 = sparse.coo_matrix((dot_vals, (rows, cols)), shape=(N, N))

# 验证两种方案结果一致
assert np.allclose(A1.todense(), A2.todense())

性能优势说明 #

• NumPy的X @ X.T调用BLAS底层优化指令，是纯C级别的运算速度，远快于Python级别的循环。
• SciPy稀疏矩阵的操作会自动跳过零元素，只处理需要计算的位置，没有冗余计算。
• 批量点积计算采用向量化操作，能充分利用CPU的SIMD并行指令，比逐个计算点积高效得多。

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