如何使用 KDTree.query_ball_tree 在 x,y 网格中查找点集答案

【问题标题】：How to find set of points in x,y grid using KDTree.query_ball_tree如何使用 KDTree.query_ball_tree 在 x,y 网格中查找点集
【发布时间】：2015-02-15 21:45:16
【问题描述】：

我在 python 中工作，我有一个 x,y 网格，它是 numpy 数组。我需要为网格中的每个点 (x1,y1) 找到与 (x1,y1) 距离为 r 的点。 Scipy 有一个函数KDTree.query_ball_tree，它接受一个 KD Tree 对象（可以从 numpy 数组构造）和一个距离 r 作为输入，但我无法理解它是如何工作的。

例如，考虑以下几点：

[(1, 1), (2, 1), (3, 1), (4, 1), (1, 2), (2, 2), (3, 2), (4, 2), (1, 3), (2, 3), (3, 3), (4, 3), (1, 4), (2, 4), (3, 4), (4, 4)]`

我想找到距离(1,1) 2 处的所有点。输出应该是：

[(1,2),(1,3),(2,1),(3,1)]

我使用 KDTree 是因为，我想避免遍历网格的 for 循环，因为网格网格是 601x90 (YxX)，如果使用 for 循环，它在时间上不是最佳的。有人可以为我提供一个示例，说明我的情况 KDTree.query_ball_tree 吗？

【问题讨论】：

我不明白，如果你有一个对称/均匀的网格，那么找到索引或多或少是微不足道的，因为元素索引与其位置有关。在这种情况下，Kd 树完全是多余的。
不，它不是对称网格。 y点和x点之间的距离不一样。 Y 距离为 0.5，X 为 3.4。上面的网格就是一个例子。实际的网格看起来有点像这样.. [(-150,-150),(-149.5,-146.6),(-149,-143.2),(-148.5,-139.8).....]...所以，我们要在这里找到的是任何点，比如说（-149，-143.2），我们要找到半径为 2 内的点。
如果 X 和 Y 距离相同，并且点是有序的，那么您仍然不需要 Kd-tree。而且你还没有指定这个KDTree.query_ball_tree函数是在哪里（在哪个模块）定义的，所以很难回答这个问题。

标签： python scipy kdtree

【解决方案1】：

如果您要查找单个点距离内的所有点，请使用scipy.spatial.KDTree.query_ball_point 而不是query_ball_tree。后者当您需要比较 sets 个点时。

import numpy as np
from scipy.spatial import KDTree

pts = np.array([(1, 1), (2, 1), (3, 1), (4, 1), (1, 2), (2, 2), (3, 2), (4, 2), (1, 3), (2, 3), (3, 3), (4, 3), (1, 4), (2, 4), (3, 4), (4, 4)])

T = KDTree(pts)
idx = T.query_ball_point([1,1],r=2)
print pts[idx]

[[1 1]
 [2 1]
 [1 2]
 [2 2]
 [1 3]
 [3 1]]

请注意，您的输出还必须包含点 (1,1)，因为它与您的目标的距离为零。

【讨论】：

【解决方案2】：

基于@Hooked 的回答，以下在具有三个坐标值的数据集中查找具有两个已知坐标的数据点。

import numpy as np
from scipy.spatial import KDTree

pts = np.array([[1, 1, 0], [2, 1, 1], [3, 1, 2], [4, 1, 3], [1, 2, 4], [2, 2, 5], [3, 2, 6], 
                [4, 2, 7], [1, 3, 8], [2, 3, 9], [3, 3, 10], [4, 3, 11], [1, 4, 12], [2, 4, 13], [3, 4, 14], [4, 4, 15]])
pts_cut=[]
pts_cut=pts[:,0:2]

T = KDTree(pts_cut)
idx = T.query_ball_point([1,1], r=2)
print(pts[idx])

【讨论】：