Question

我在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，因为我想避免遍历网格的循环，因为网格网格是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)，因为它与您的目标距离为零。

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

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

2 个答案: