如何优化计算归一化汉明距离的高耗时numpy操作

核心瓶颈说明 #

原代码使用循环逐行计算，时间复杂度为O(N²D)，其中N=76187、D=247*20=4940，重复的向量点积计算和Python循环开销是耗时过高的核心原因。

优化方案 #

1. 全量向量化计算（内存足够时首选，逻辑完全等价） #

直接利用numpy的矩阵运算替代Python循环，底层调用高度优化的BLAS库，可将耗时降低两个数量级以上：

import numpy as np

cutoff = 0.2
# 展平输入
X_flat = X.reshape(X.shape[0], -1)
# 预先计算每个向量的模方，避免重复计算
x_sq = np.sum(X_flat ** 2, axis=1)
threshold = 1 - cutoff
# 一次性计算所有两两向量的点积
dot_mat = X_flat @ X_flat.T
# 批量判断并统计符合条件的数量
N_list = 1.0 / np.sum(dot_mat > threshold * x_sq[:, None], axis=1)

2. 分块计算（内存不足时使用） #

全量76187*76187的点积矩阵如果用float64存储约占44GB内存，内存不足时可采用分块计算方案，内存占用可控制在GB级别：

block_size = 1024 # 可根据可用内存大小调整
N = X_flat.shape[0]
N_list = np.zeros(N, dtype=np.float64)
threshold = 1 - cutoff

for i in range(0, N, block_size):
    end_idx = min(i + block_size, N)
    # 只计算当前块的点积
    block_dot = X_flat[i:end_idx] @ X_flat.T
    block_threshold = threshold * x_sq[i:end_idx, None]
    N_list[i:end_idx] = 1.0 / np.sum(block_dot > block_threshold, axis=1)

3. 进一步加速方案 #

• 精度要求不高时，将X_flat转换为float32类型计算，速度可提升1倍以上，内存占用减半
• 有GPU可用时，可将运算迁移到PyTorch/TensorFlow框架执行，7万量级的两两计算仅需数秒到数十秒即可完成
• 可使用Numba对分块逻辑做JIT编译，进一步降低循环开销

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