使用KDTree.query_ball_tree如何在x，y网格中找到点集

Question

我正在使用Python编程，有两个numpy数组x和y组成的网格。我需要找到距离(x1,y1)点r距离的点。Scipy提供了一个函数KDTree.query_ball_tree，该函数需要一个KD树对象（可以从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函数吗？

- Numby

---

我不明白，如果您有一个对称/均匀的网格，那么找到索引或多或少是微不足道的，因为元素索引与其位置有关。在这种情况下，Kd树完全是杀鸡焉用牛刀。 - Hannes Ovrén

不，它不是对称网格。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内的点。 - Numby

如果X和Y距离相同，并且点是有序的，那么您仍然不需要Kd树。您还没有指定KDTree.query_ball_tree函数在哪个模块中定义，这使得回答问题变得困难。 - Hannes Ovrén

Answer 1

如果您想查找与单个点距离在一定范围内的所有点，请使用scipy.spatial.KDTree.query_ball_point而不是query_ball_tree。后者用于比较彼此之间的点集。

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） ，因为它是从目标点距离为零。

- Hooked

Answer 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])

- U3.1415926