Python高效查找n个最近邻向量：词嵌入相似度计算优化问询

嘿，我太懂你现在的困扰了——两百万个200维向量挨个算余弦相似度，完全是暴力解法，时间复杂度O(n)，数据量上来之后速度肯定崩。咱换点高效的思路，下面几个方案从易到难，你可以按需选：

1. 先归一化向量 + 矩阵批量运算（快速优化现有逻辑） #

余弦相似度的计算可以转化为归一化向量的点积（因为cosθ = (A·B)/(|A||B|)，如果A和B都归一化到模长1，那直接算点积就是余弦相似度）。用numpy的矩阵运算代替Python循环，能把速度提升好几个数量级。

先预处理你的词嵌入向量（只做一次就行，后续查询直接复用）：

import numpy as np

def normalize_vectors(word_vectors):
    vectors = np.array(list(word_vectors.values()))
    norms = np.linalg.norm(vectors, axis=1, keepdims=True)
    normalized_vectors = vectors / norms
    word_list = list(word_vectors.keys())
    vec_matrix = normalized_vectors
    # 返回单词列表、归一化向量矩阵、单词-向量映射
    return word_list, vec_matrix, {word: vec for word, vec in zip(word_list, normalized_vectors)}

然后修改最近邻查询函数：

def n_nearest_words(word, n, word_list, vec_matrix):
    """
    基于预归一化的向量矩阵快速计算最近邻
    :param word: 目标单词
    :param n: 需要返回的最近邻数量
    :param word_list: 所有单词的列表（和vec_matrix的行一一对应）
    :param vec_matrix: 所有归一化向量组成的numpy矩阵（每行是一个单词的向量）
    """
    target_idx = word_list.index(word)
    target_vec = vec_matrix[target_idx]
    
    # 批量计算所有向量与目标向量的点积（即余弦相似度）
    similarities = np.dot(vec_matrix, target_vec)
    
    # 取相似度最高的n+1个（去掉目标单词自己）
    top_indices = np.argsort(similarities)[::-1][1:n+1]
    
    return [word_list[idx] for idx in top_indices]

2. 用近似最近邻(ANN)库（超大规模数据的最优解） #

如果两百万个向量还是觉得慢，就得用专门的近似最近邻库，比如FAISS（Facebook开源，适合大规模数据）或者Annoy（Spotify开源，轻量易用）。这些库会给向量建立索引，查询时间复杂度降到O(log n)，速度提升非常明显。

示例：用FAISS实现 #

import faiss
import numpy as np

def build_faiss_index(word_vectors):
    word_list = list(word_vectors.keys())
    vectors = np.array(list(word_vectors.values())).astype('float32')
    faiss.normalize_L2(vectors)
    
    d = vectors.shape[1]  # 向量维度（这里是200）
    nlist = 1000  # 聚类中心数量，可根据数据量调整
    # 用内积度量（对应归一化向量的余弦相似度）
    index = faiss.IndexIVFFlat(faiss.IndexFlatIP(d), d, nlist, faiss.METRIC_INNER_PRODUCT)
    index.train(vectors)
    index.add(vectors)
    
    return word_list, index

def n_nearest_words_faiss(word, n, word_list, index, word_vectors):
    target_vec = np.array([word_vectors[word]]).astype('float32')
    faiss.normalize_L2(target_vec)
    
    # 查询n+1个最近邻，过滤掉目标单词自己
    distances, indices = index.search(target_vec, n+1)
    return [word_list[idx] for idx in indices[0] if word_list[idx] != word][:n]

示例：用Annoy实现 #

from annoy import AnnoyIndex

def build_annoy_index(word_vectors):
    d = 200  # 向量维度
    # angular度量对应余弦相似度
    index = AnnoyIndex(d, metric='angular')
    
    word_list = list(word_vectors.keys())
    for i, word in enumerate(word_list):
        index.add_item(i, word_vectors[word])
    
    index.build(10)  # 树的数量，值越大精度越高，速度稍慢
    return word_list, index

def n_nearest_words_annoy(word, n, word_list, index):
    target_idx = word_list.index(word)
    # 查询n+1个最近邻，过滤掉目标单词自己
    nearest_indices = index.get_nns_by_item(target_idx, n+1)
    return [word_list[idx] for idx in nearest_indices if idx != target_idx][:n]

3. 额外小技巧 #

• 频繁查询时，一定要提前建好索引/归一化矩阵，不要每次查询都重复处理
• 如果有GPU，FAISS支持GPU加速，速度还能再上一个台阶
• 可以把词向量和单词列表保存到磁盘（比如用numpy的save/load），避免每次重启都重新处理

你原本的暴力解法代码：

def n_nearest_words(word, n, word_vectors):
    """ 
    Return a list of the n nearest words to param word, based on cosine similarity
    param wo...
    """
    # 逐一计算余弦相似度的慢逻辑

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