多权重场景下无需重复调用KNN实现向量相似度重排序

我现有一段用于计算存储在values变量中的多向量间最近邻的代码，这些向量的取值依赖权重参数，每次迭代时向量的每一列都会对应不同的权重。
示例代码每次都会查找最后一个向量（vector[3]）的最近邻，极简实现如下：

from sklearn.neighbors import NearestNeighbors
knn = NearestNeighbors(n_neighbors=1)

values = [
    [2, 5, 1],
    [4, 2, 3],
    [1, 5, 2],
    [4, 5, 4]
]

weights = [
    [1, 3, 1],
    [0.5, 2, 1],
    [3, 1, 2]
]

# weights set No1
new_values = []
for line in values:
    new_values.append([a*b for a,b in zip(line,weights[0])])    
knn.fit(new_values)
print(knn.kneighbors(new_values[3]))

# weights set No2
new_values = []
for line in values:
    new_values.append([a*b for a,b in zip(line,weights[1])])         
knn.fit(new_values)
print(knn.kneighbors(new_values[3]))

# weights set No3
new_values = []
for line in values:
    new_values.append([a*b for a,b in zip(line,weights[2])])        
knn.fit(new_values)
print(knn.kneighbors(new_values[3]))

注：上述写法仅为展示逻辑重复问题，实际场景中可以通过for循环遍历不同权重集。

核心疑问：

• 是否可以避免多次重复调用KNN，仅在初始阶段调用一次KNN完成初始相似度排名/排序，后续仅通过重计算得到各权重下的结果？
• 是否可以通过减少KNN的调用次数，降低这段代码的计算复杂度？

补充说明：已知存在比Scikit-Learn实现速度更快的KNN版本，但这不是问题核心，重点是实现仅调用1次KNN，而非示例中N次重复调用的逻辑。

首先明确结论：任意权重变化的场景下，你无法通过单次无权重KNN的预计算排序结果直接得到所有权重集下的准确最近邻——因为列权重的极端变化可能让原本距离很远的点变成最近邻，初始排序没有参考性。但你完全可以通过拆解加权距离的计算逻辑，避免每次重建KNN索引的冗余开销，实现比重复调用KNN低得多的计算成本，具体逻辑如下：

核心原理 #

Scikit-Learn中KNN默认使用的欧氏距离，在逐列加权场景下可以拆分为固定项和可变项两部分：
两个向量x和y的加权欧氏距离平方为：
$$d_w(x,y)^2 = \sum_{i=1}^d w_i (x_i - y_i)^2$$
其中$d$是向量维度，$w_i$是第i列的权重，$(x_i-y_i)^2$是两个向量在第i列的原始差值平方——这部分和权重完全无关，是可以提前一次性计算的固定值。

优化实现方案 #

你只需要提前预计算所有候选向量和查询向量（示例中是最后一个向量values[3]）在每个维度上的差值平方矩阵，后续更换权重时不需要重新拟合KNN，仅做一次矩阵向量乘法就能得到所有距离，直接取最小值即可，计算量远低于重复拟合KNN：

import numpy as np

values = np.array([
    [2, 5, 1],
    [4, 2, 3],
    [1, 5, 2],
    [4, 5, 4]
])
weights = np.array([
    [1, 3, 1],
    [0.5, 2, 1],
    [3, 1, 2]
])

# 仅需执行1次的预计算：提前算所有候选向量和查询向量的逐维差值平方
query_vec = values[-1]
# 排除查询向量本身，计算其余向量和它的逐维差平方
dim_diff_sq = (values[:-1] - query_vec) ** 2

# 后续遍历所有权重集，不需要再调用KNN
for w in weights:
    # 矩阵乘法直接得到所有候选向量到查询点的加权距离平方
    dists = dim_diff_sq @ w
    # 取距离最小的索引就是最近邻，和原KNN输出结果完全一致
    nearest_idx = dists.argmin()
    nearest_dist = np.sqrt(dists[nearest_idx])
    print(f"最近邻索引：{nearest_idx}，距离：{nearest_dist}")

方案说明 #

• 该方案和原代码输出结果完全一致，无精度损失
• 预计算阶段仅需执行一次逐维差值计算，后续每个权重集的计算仅为低复杂度的矩阵向量乘法，时间复杂度远低于重复拟合KNN、重复构建空间索引的开销
• 如果需要查询的不是单个固定向量，而是所有向量的K近邻，只需要提前预计算所有向量对之间的逐维差值平方张量即可，后续更换权重时同样仅需张量和权重的乘法就能得到全量距离矩阵，不需要重复调用KNN
• 如果后续权重变化范围很小，还可以在预计算阶段提前筛选距离目标点较近的候选集，进一步缩小后续计算范围；但如果权重可能出现极端值（比如某列权重远大于其他列），候选集需要覆盖足够多的近邻才能保证结果准确。

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