如何用 Python 实现 KD-Tree 并在高维超立方体点集中执行最近邻搜索

如何用 Python 实现 KD-Tree 并在高维超立方体点集中执行最近邻搜索

【免费下载链接】PythonAll Algorithms implemented in Python项目地址: https://gitcode.com/GitHub_Trending/pyt/Python

如果你的任务是在 N 维空间中对一组已知点反复做"给定查询点，找出最近的点"这类搜索，Python 仓库 data_structures/kd_tree 目录下提供了一条可以直接运行的实现路径：用 build_kdtree.py 从点集构建树，用 nearest_neighbour_search.py 执行最近邻搜索，并用 example/example_usage.py 在 10 维超立方体上跑一个完整演示。整个过程只需要 NumPy 一个第三方依赖（pyproject.toml 中固定为numpy>=2.1.3），Python 版本要求为 3.14 及以上（requires-python = ">=3.14"）。

模块组成：三个核心文件各承担什么

该实现刻意把职责拆成三个文件，读代码时按这个顺序看即可：

建树的关键逻辑在build_kdtree中：

k = len(points[0]) # Dimensionality of the points axis = depth % k # 当前层沿哪个维度切分 # Sort point list and choose median as pivot element points.sort(key=lambda point: point[axis]) median_idx = len(points) // 2 left_points = points[:median_idx] right_points = points[median_idx + 1 :]

即按depth % k轮换切分轴，按该轴排序后取中位数点作为节点，左右子集递归建树。注意points.sort(...)直接修改传入列表的顺序，演示代码里传的是points.tolist()生成的新列表，如果你传自己的列表，建完树后不要假设原顺序不变。

搜索的关键逻辑在nearest_neighbour_search中：先递归进入查询点所在的近侧子树，再用当前轴上的切分距离判断远侧子树是否可能包含更近的点：

# If the further subtree has a closer point if (query_point[axis] - current_point[axis]) ** 2 < nearest_dist: search(further_subtree, depth + 1)

函数的返回契约（来自函数 docstring 和测试用例）：

• nearest_point：离查询点最近的点；树为空时是None；
• nearest_dist：到最近点的平方距离（不是欧氏距离本身），树为空时是float("inf")；
• nodes_visited：搜索过程中访问的节点数，树为空时是0。

需要距离本身时自行对nearest_dist开平方即可。

执行步骤：运行 10 维超立方体演示

example_usage.py 的main()是仓库给出的完整用法，其参数选择为：

• num_points = 5000：点数量；
• cube_size = 10.0：超立方体边长；
• num_dimensions = 10：维度。

点由 example/hypercube_points.py 中的hypercube_points(num_points, hypercube_size, num_dimensions)生成，它是用np.random.default_rng()在[0, hypercube_size)内均匀采样，返回形状(num_points, num_dimensions)的数组：

rng = np.random.default_rng() shape = (num_points, num_dimensions) return hypercube_size * rng.random(shape)

完整流程就是：

points: np.ndarray = hypercube_points(num_points, cube_size, num_dimensions) hypercube_kdtree = build_kdtree(points.tolist()) # Generate a random query point within the same space rng = np.random.default_rng() query_point: list[float] = rng.random(num_dimensions).tolist() # Perform nearest neighbor search nearest_point, nearest_dist, nodes_visited = nearest_neighbour_search( hypercube_kdtree, query_point )

两点值得注意：一是build_kdtree接收list[list[float]]，所以 numpy 数组要先.tolist()；二是演示中的查询点由rng.random(num_dimensions)生成，各维落在[0, 1)区间，位于边长为 10 的超立方体内部，这是文档代码给出的取样方式，不是硬性要求。

在仓库根目录下运行（导入语句是from data_structures.kd_tree...这种包路径，必须在仓库根目录执行）：

python data_structures/kd_tree/example/example_usage.py

前提是当前 Python 环境已安装 NumPy（项目固定版本为numpy>=2.1.3，Python>=3.14）。

结果验证：输出解读与测试用例

演示程序的输出格式固定为四行（示例结果，数值随随机点变化）：

Query point: [...] Nearest point: [...] Distance: 0.xxxx Nodes visited: N

其中Distance打印的是平方距离（源码中用:.4f格式化）。由于每次运行点集和查询点都随机生成，不要把这些数值当成固定预期；判断运行正常的标准是：Nearest point不为None、Distance非负、Nodes visited为 0 以上的整数。

要稳定地验证实现本身，跑目录下的 pytest 测试 tests/test_kdtree.py：

python -m pytest data_structures/kd_tree/tests/test_kdtree.py

该测试文件覆盖了三类断言，可直接作为核对依据：

1. 建树（test_build_kdtree，参数化三个用例）：
   • 空点列表（num_points=0）时build_kdtree必须返回None；
   • 2 维 10 点、depth=2，以及 3 维 10 点、depth=-2（负深度按depth % k处理，不报错）时，返回的根节点必须是KDNode，且len(kdtree.point) == num_dimensions。
2. 搜索（test_nearest_neighbour_search）：2 维、10 个超立方体点（cube_size=10.0）建树后，用rng.random(2)生成查询点，断言nearest_point is not None、nearest_dist >= 0、nodes_visited >= 0。
3. 边界（test_edge_cases）：对空树build_kdtree([])搜索 2 维查询点[0.0, 0.0]，断言返回(None, float("inf"), 0)——这是空树行为最直接的判定标准。

如果你的点集很小，也可以照test_nearest_neighbour_search的写法把num_points、num_dimensions换小（如 2 维 10 点）再跑一遍，输出应满足与测试相同的三个断言。

限制与适用边界

结合源码和测试，实现文档明确给出的边界是：

• 距离度量是平方欧氏距离：搜索比较和返回的都是平方值，与查询结果做距离比较时要保持口径一致；
• 空树有明确返回值（None/inf/0），不需要额外判空即可处理，但查询点维数必须与树中点一致（搜索内层用zip(query_point, current_point)配对计算，维数不一致会导致距离算错而不是报错）；
• 切分轴只由depth % k决定，k取自第一个点的维度（len(points[0])），因此点集内各点维数应保持一致；
• 演示的 10 维 5000 点场景说明该实现面向"中高维"场景；仓库 README 也注明所有实现均以教学为目的，可能与标准库或专用库的效率不同，性能敏感场景需自行评估。

跑通演示与测试后，这套代码可以作为独立模块继续使用：把build_kdtree和nearest_neighbour_search两个函数按list[list[float]]契约接入你自己的点集即可。

【免费下载链接】PythonAll Algorithms implemented in Python项目地址: https://gitcode.com/GitHub_Trending/pyt/Python