首页 > 解决方案 > 了解 scipy.spatial.KDTree 中的 `leafsize`

问题描述

问题陈述:

我在 3D 空间中有 150k 个点,它们的坐标存储在尺寸为 [150k, 3] 的矩阵中,以 mm 为单位。

我想找到p半径内给定点的所有邻居r。我想以最准确的方式做到这一点。

我应该如何选择我的leafsize参数?

from scipy.spatial import KDTree
import numpy as np

pts = np.random.rand(150000,3)

T1 = KDTree(pts, leafsize=20)
T2 = KDTree(pts, leafsize=1)

neighbors1= T1.query_ball_point((0.3,0.2,0.1), r=2.0)
neighbors2= T2.query_ball_point((0.3,0.2,0.1), r=2.0)

np.allclose(sorted(neighbors1), sorted(neighbors2))
True

标签: scipytreenearest-neighborkdtree

解决方案


该函数query_ball_point将为任何版本的搜索树返回正确的点集。该leafsize参数不影响查询的结果,只影响结果的性能。

想象一下下面显示的两棵树用于相同的数据(但不同的叶子大小参数)和一个搜索红色圆圈内所有点的查询。

示例搜索树

在这两种情况下,代码都只会返回红圈内的两个点。这是通过检查与圆相交的树的所有框中的所有点来完成的。这导致每种情况下的工作量不同(即不同的性能)。对于左树(对应更大的叶子大小),算法必须检查圆内是否有 13 个点(上相交框中有 6 个,下相交框中有 7 个)。在右树(叶子尺寸较小)中,只处理了三个点(一个在上相交框中,两个在下相交框中)。

按照这个逻辑,您可能认为总是使用较小的叶子大小才有意义:这将最大限度地减少算法结束时的实际比较次数(确定点是否实际位于查询区域中)。但这并不是那么简单:较小的叶子大小将生成更深的树,从而增加了构建时间和树遍历时间的成本。在树遍历性能与叶级比较之间取得适当的平衡实际上取决于进入树的数据类型以及您正在进行的特定叶级比较。这就是为什么 scipy 提供 Leafsize 参数作为参数,以便您可以调整事物以在特定算法上表现最佳。


推荐阅读