如何在scipy KDTree中高效查找某位置两半径区间内的点

要高效查找距离指定点满足 r1 < 距离 < r2 的点，一次范围查询加本地过滤是比两次query_ball_point加差集更优的方案，具体思路和实现如下：

核心思路 #

先通过KDTree的query_ball_point获取所有距离目标点小于等于r2的点，再在这些点中过滤掉距离小于等于r1的点。这种方法只做一次KDTree的范围查询，后续用numpy向量运算做过滤，整体开销远低于两次查询加集合差集。

代码实现 #

import numpy as np
from scipy.spatial import KDTree

# 示例数据：10000个3维随机点
data = np.random.rand(10000, 3)
kd_tree = KDTree(data)
target = np.array([0.5, 0.5, 0.5])
r1 = 0.1
r2 = 0.3

# 步骤1：获取r2范围内的所有点索引（关闭排序减少开销）
indices_r2 = kd_tree.query_ball_point(target, r2, return_sorted=False)
# 步骤2：取出对应点的坐标
points_r2 = data[indices_r2]
# 步骤3：批量计算这些点到目标点的距离
distances = np.linalg.norm(points_r2 - target, axis=1)
# 步骤4：筛选出距离大于r1的点索引
filtered_indices = np.array(indices_r2)[distances > r1]

为什么这个方案更优 #

1. 减少KDTree查询次数：KDTree的范围查询时间复杂度是O(log n + k)（n是总点数，k是结果点数），一次查询比两次查询节省了一次log n的开销，以及一次小范围查询的k1（r1内的点数）处理成本。
2. 过滤效率更高：numpy的向量运算是C级别的批量计算，比Python原生列表转集合做差集的操作快得多，尤其是当r2范围内的点数量较大时，优势更明显。
3. 可选优化：添加return_sorted=False参数可以避免KDTree对结果排序，进一步减少不必要的开销（因为后续过滤不需要有序的索引）。

原方案的劣势 #

两次query_ball_point加集合差集的做法，在数据量较大时会面临两个问题：

• 两次KDTree查询的累积开销；
• 大列表转集合、求差集的Python层面循环开销，远慢于numpy的批量运算。

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