如何高效计算仅保留C中非零元素的Sparse × Sparse矩阵乘积

嘿，这个需求我之前也碰到过——用密集矩阵先乘再清零的方式确实太浪费资源，尤其是矩阵规模大的时候，完全是在做无用功。既然内存不受限制，咱们直接利用稀疏矩阵的特性来优化就行，下面给你两个实用的方案：

方案一：利用Scipy稀疏矩阵的元素级乘法（最通用） #

Scipy的sparse模块对稀疏矩阵的操作做了极致优化，咱们可以先计算A·B的稀疏矩阵结果，再和C的非零掩码做元素级乘法，这样只会保留两者都非零的位置（正好是你要的C中非零的位置）。

代码示例：

import numpy as np
import scipy.sparse as sp

# 假设A、B、C都是scipy的稀疏矩阵（比如csr_matrix格式）
# 先转成CSR格式，乘法效率最高
A = A.tocsr()
B = B.tocsr()
C = C.tocsr()

# 计算A和B的乘积（稀疏矩阵乘法，只处理非零元素）
prod_sparse = A.dot(B)

# 创建C的二进制掩码矩阵：非零位置为1，零位置为0
mask = sp.csr_matrix(
    (np.ones_like(C.data), C.indices, C.indptr),
    shape=C.shape
)

# 只保留C中非零位置的乘积结果
result = prod_sparse.multiply(mask)

这个方法的优势在于：所有操作都在稀疏矩阵上完成，不会生成庞大的密集矩阵，计算量只和非零元素的数量挂钩，比原方法快几个数量级都很正常。

方案二：仅计算C中非零位置的点积（极端稀疏场景更优） #

如果C的非零元素数量远少于A·B的非零元素数量，那咱们可以更进一步——直接遍历C的每个非零位置，计算对应A的行和B的列的点积，完全跳过不需要的位置。

代码示例：

import numpy as np
import scipy.sparse as sp

def sparse_prod_filtered(A, B, C):
    # 统一转成CSR格式方便操作
    A = A.tocsr()
    B = B.tocsr()
    C = C.tocsr()
    
    # 获取C所有非零元素的行和列索引
    rows, cols = C.nonzero()
    values = []
    
    # 逐个计算指定位置的乘积值
    for i, j in zip(rows, cols):
        # 计算A第i行和B第j列的点积
        dot_val = A[i].dot(B[:, j]).toarray()[0, 0]
        values.append(dot_val)
    
    # 用计算出的值构建结果稀疏矩阵
    result = sp.csr_matrix((values, (rows, cols)), shape=C.shape)
    return result

这个方法适合C非常稀疏的场景，比如C的非零元素只有几百个，但A·B有几百万个非零元素——此时直接计算需要的点积，比先算整个乘积再过滤要高效得多。

关键注意事项 #

• 优先使用CSR/CSC格式：Scipy的稀疏矩阵乘法在这两种格式下的优化最好，其他格式（比如COO）可以先转成CSR再操作。
• 避免密集化：绝对不要把稀疏矩阵转成numpy.array来操作，哪怕内存够，计算速度也会慢到离谱。

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