Python中加速稀疏成对向量距离计算的优化方法问询

核心思路：仅计算mask标记的目标位置 #

既然只需要保留M中1对应的向量对距离，直接定位这些位置，单独计算对应向量对的距离，完全避免生成全量距离矩阵的冗余开销。

方法一：基于Numpy索引定位的轻量实现 #

利用Numpy的坐标索引提取需要计算的向量对，批量计算距离：

import numpy as np

A = np.array([[1, 2], [2, 3], [3, 4]])
B = np.array([[4, 5], [5, 6], [6, 7], [7, 8], [8, 9]])
M = np.array([[0, 0, 0, 1, 0], [1, 1, 0, 0, 0], [0, 0, 0, 0, 1]])

# 获取M中1的坐标位置
rows, cols = np.where(M == 1)
# 提取对应位置的向量对
a_target = A[rows]
b_target = B[cols]
# 批量计算欧氏距离
distances = np.linalg.norm(a_target - b_target, ord=2, axis=1)
# 可选：用Scipy稀疏矩阵存储结果，节省内存
from scipy.sparse import csr_matrix
masked_distances = csr_matrix((distances, (rows, cols)), shape=M.shape)

这种方式的计算量仅为需要保留的向量对数量，当占比<0.01时，计算量直接降至原方案的1%以内，内存占用也大幅降低。

方法二：基于Scipy稀疏矩阵的原生处理 #

如果M本身是稀疏矩阵格式，直接利用稀疏矩阵的非零坐标信息计算：

from scipy.sparse import csr_matrix

# 假设M是csr格式（大规模场景下推荐用稀疏格式存储M）
M_csr = csr_matrix(M)
rows, cols = M_csr.nonzero()
# 提取目标向量对并计算距离
a_target = A[rows]
b_target = B[cols]
distances = np.linalg.norm(a_target - b_target, ord=2, axis=1)
# 构建稀疏结果矩阵
result = csr_matrix((distances, (rows, cols)), shape=M_csr.shape)

避免了稀疏矩阵转密集矩阵的额外开销，适合超大规模向量组场景。

方法三：用FAISS加速超大规模向量计算 #

如果向量数量达十万级以上、维度较高，Facebook的FAISS库通过底层C++优化，能高效完成这类稀疏距离计算：

import faiss
import numpy as np

# 将向量转为FAISS支持的float32格式
A_float = A.astype(np.float32)
B_float = B.astype(np.float32)

# 构建FAISS索引
index = faiss.IndexFlatL2(B.shape[1])
index.add(B_float)

# 逐行处理A，只查询M标记的B向量
distances = []
row_indices = []
col_indices = []

for i in range(A.shape[0]):
    target_cols = np.where(M[i] == 1)[0]
    if len(target_cols) == 0:
        continue
    # 单独查询目标向量的距离
    D, _ = index.search(A_float[i:i+1], len(target_cols))
    distances.extend(D[0])
    row_indices.extend([i]*len(target_cols))
    col_indices.extend(target_cols)

# 构建稀疏结果矩阵
result = csr_matrix((distances, (row_indices, col_indices)), shape=M.shape)

FAISS的检索效率远高于纯Python实现，适合超大规模向量场景。

原方案效率低的原因 #

• 原代码生成全量diff矩阵（形状为(N,M,D)），当N、M为数千级时，计算量和内存占用均为O(N*M)，冗余极高；
• np.vectorize本质是Python循环封装，批量计算效率远低于直接索引的向量运算；
• 高维稀疏数组操作复杂度高，不如直接提取非零坐标进行针对性计算。

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