如何使用KDTree.query球在x，y网格中查找点集_

Question

我在python中工作，我有一个x，y网格，它是numpy数组。我需要为网格中的每个点（x1，y1）找到距离（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循环，它在时间上不会是最佳的。有人能给我举个例子来说明我的情况吗？