Sklearn近邻算法运行时疑问：模型构建与查询耗时差异及优化

问题背景 #

我有一个尺寸约为4000×10的2D numpy数组，将每行视为10维空间中的点，希望计算每个点的5近邻。运行了以下代码：

import numpy as np
from sklearn.neighbors import NearestNeighbors
import time
A = np.random.rand(4000, 10)
t1 = time.time()
nbrs = NearestNeighbors(n_neighbors = 5, algorithm = 'kd_tree').fit(A)
t2 = time.time()
distance, indices = nbrs.kneighbors(A)
t3 = time.time()
print('time taken for step1: fitting is:', t2 - t1)
print('time taken for step2: retrieving data is:', t3 - t2)

运行结果：

time taken for step1: fitting is: 0.009654521942138672
time taken for step2: retrieving data is: 0.3108406066894531

问题一：为何获取距离/索引的耗时远高于拟合过程？ #

你的理解搞反了两个阶段的核心工作：

• fit()阶段只是构建kd-tree的索引结构，仅需把数据集按维度划分、生成树的节点，全程不需要计算点之间的距离，计算量极小，所以耗时短。
• kneighbors()阶段才是真正的近邻搜索过程：要对4000个点逐个执行搜索，每个点都得遍历kd-tree，计算与候选点的距离，筛选出最近的5个。这是批量执行的高计算量操作，自然耗时远高于树的构建。

问题二：仅需近邻索引时如何加快速度？ #

可以从这几个方向优化：

• 跳过距离计算：调用kneighbors()时传入return_distance=False，函数只会返回索引数组，省去距离计算和存储的开销，代码示例：

  indices = nbrs.kneighbors(A, return_distance=False)
  

• 换用更适配的算法：10维属于中等维度，ball_tree在这类场景下的搜索效率可能优于kd_tree；也可以尝试brute暴力搜索（虽然高维通常暴力慢，但4000个点+10维的规模可以测试对比），修改算法参数即可：

  nbrs = NearestNeighbors(n_neighbors=5, algorithm='ball_tree').fit(A)
  

• 调整kd-tree的叶子节点大小：修改leaf_size参数（默认30），更小的leaf_size会让树结构更精细，构建时间变长但搜索速度可能提升；反之更大的leaf_size会让构建更快但搜索变慢，你可以根据自己的数据集测试最优值。

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