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

Question

我在python中工作,我有ax,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)]`

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

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

我正在使用KDTree,因为我想避免遍历网格的循环,因为网格网格是601x90(YxX),如果使用循环,它将不是最佳的时间.有人能为我提供一个说明KDTree.query_ball_tree我情况的例子吗？

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),因为它与目标的距离为零.