Scipy KDTree获取由两点定义的网格的矩形子集

时间:2019-01-22 15:41:23

标签: python scipy kdtree

我正在使用以下示例:

from scipy import spatial
x, y = np.mgrid[0:5, 2:8]
tree = spatial.KDTree(list(zip(x.ravel(), y.ravel())))
pts = np.array([[0, 0], [2.1, 2.9]])
idx = tree.query(pts)[1]
data = tree.data[??????????]

如果我输入两个任意点(请参见变量pts),则希望返回位于由两个点定义的矩形内的所有坐标对(KDTree查找最近的邻居)。因此,在这种情况下:

array([[0, 0],
       [0, 1],
       [0, 2],
       [1, 0],
       [1, 1],
       [1, 2],
       [2, 0],
       [2, 1],
       [2, 2]])

如何从树数据中实现这一目标?

1 个答案:

答案 0 :(得分:0)

似乎我找到了解决方法:

from scipy import spatial
import numpy as np
x, y = np.mgrid[0:5, 0:5]
tree = spatial.KDTree(list(zip(x.ravel(), y.ravel())))
pts = np.array([[0, 0], [2.1, 2.2]])
idx = tree.query(pts)[1]
data = tree.data[[idx[0], idx[1]]]
rectangle = tree.data[np.where((tree.data[:,0]>=min(data[:,0])) & (tree.data[:,0]<=max(data[:,0])) & (tree.data[:,1]>=min(data[:,1])) & (tree.data[:,1]<=max(data[:,1])))]

但是,我很想看到使用查询选项的解决方案!