使用 k-d 树搜索相邻点答案

【问题标题】：Search neighboring points using k-d trees使用 k-d 树搜索相邻点
【发布时间】：2016-05-03 08:48:59
【问题描述】：

我有一堆坐标包含在三个numpy 数组中：xarr、yarr 和 zarr（每个数组中的对应位置属于同一个点 - 即第一个点在 xarr[0]， yarr[0]，zarr[0])。给定空间中的另一个点P(x,y,z)，我想找到距离P(x,y,z) r 内的所有点。

我目前（而且效率非常低）的方法是简单地迭代并计算到每个点的距离，看看它是否在P(x,y,z) 的 r 内。

但是，我想使用 SciPy 的 k-d 树算法来执行此操作，但我不确定如何开始实现它（我对 Python 很陌生）。如果有人能简要地概述一些代码来演示如何设置给定数据的 k-d 树，我将不胜感激。

我知道SciPy documentation of its k-d tree implementation，我已经看过了，但我仍然对如何根据我拥有的格式的数据创建树感到困惑（np.mgrid 和 ravel() 被调用，我不太明白为什么）。

谢谢！

【问题讨论】：

如果我理解正确，您有一个 single Nx3 点数组，对于该数组中的每个点，您要计算在某个半径内的其他点的数量呢？
@ali_m 他们有3个独立的数组，每个坐标一个，他们想在所描述的点中找到靠近新点的点。
我不知道这对这个特定问题是否有用，我对 k-d 树并不感兴趣（尽管我可能应该），但我发现这是一个有用的库：@ 987654322@

标签： python numpy scipy binary-search-tree

【解决方案1】：

省略导入...我没有你的数据，所以我必须伪造一些...

In [40]: x = np.random.random(100)
In [41]: y = np.random.random(100)    
In [42]: z = np.random.random(100)    
In [43]: p = np.random.random(3)    
In [44]: p
Out[44]: array([ 0.60515083,  0.39263392,  0.36129813])

即三个坐标数组和一个我将搜索邻居的点。

接下来，让我们看看如何构造一个包含与不同数据点和三个列一样多的行的数组...

In [45]: np.vstack((x,y,z)).T.shape
Out[45]: (100, 3)

嗯，没错。

我们使用来自scipy.spatial 的KDTree 构建k-d 树

In [46]: tree = KDTree(np.vstack((x,y,z)).T)

然后我们使用树的一种方法，恰如其分地命名为.query_ball_point()，来查找靠近p的点的索引

In [47]: indices = tree.query_ball_point(p, 0.33)

我任意使用的半径等于 1/3。

最终我们希望看到这些邻居，所以我将使用树的 .data 属性和我刚刚计算的索引

In [48]: tree.data[indices]
Out[48]: 
array([[ 0.4117843 ,  0.21440852,  0.3352732 ],
       [ 0.48921727,  0.13855976,  0.43331816],
       [ 0.71598133,  0.32270361,  0.20292187],
       [ 0.71761991,  0.27309708,  0.12670474],
       [ 0.6282775 ,  0.13752325,  0.4143872 ],
       [ 0.55995847,  0.31302848,  0.2780926 ],
       [ 0.75896359,  0.16043536,  0.33530071],
       [ 0.81138529,  0.64635994,  0.33819097],
       [ 0.43537193,  0.5353203 ,  0.52095431],
       [ 0.66996807,  0.48346547,  0.52761835],
       [ 0.69426851,  0.24725511,  0.57650329],
       [ 0.5350322 ,  0.23155768,  0.62545958],
       [ 0.51228139,  0.38078056,  0.61246054]])

仅此而已……

【讨论】：

这太棒了！非常感谢，这正是我想要的。

【解决方案2】：

以下是 scipy 文档中提供的示例的解释：

from scipy import spatial
x, y = np.mgrid[0:4, 0:4]

np.mgrid 创建一个从 0 到 4 的网格 x,y。由于您已经有了 x,y,z 坐标，因此您将跳过这一步。

points = zip(x.ravel(), y.ravel())
points = zip(xarr.ravel(), yarr.ravel(), zarr.ravel()) #in your case
points = zip(xarr, yarr, zarr) # if x,y,z are already 1-d

zip 创建一个包含每个 x,y 点对的元组列表（将每个点的坐标关联在一起）。 ravel 展平 x、y 网格（将 n-d 数组转换为 1-d），以便可以使用 zip。在您的情况下，如果xarr、yarr、zarr 还不是一维，您将只使用ravel。

tree = spatial.KDTree(points)

创建索引点以提供快速邻居查找。

tree.query_ball_point([2, 0], 1)

在点[2,0]的r=1内查找点

希望这会有所帮助。

【讨论】：