Python中KDTree的点索引

Question

给定一组点，如何在KDTree中获取它们的索引？

from scipy import spatial
import numpy as np

#some data
x, y = np.mgrid[0:3, 0:3]
data = zip(x.ravel(), y.ravel())

points = [[0,1], [2,2]]

#KDTree
tree = spatial.cKDTree(data)

# incices of points in tree should be [1,8]

我可以做类似这样的操作：

[tree.query_ball_point(i,r=0) for i in points]

>>> [[1], [8]]

用这种方式做有意义吗？

- a.smiet

Answer 1

使用cKDTree.query(x, k, ...)函数，查找给定点集x的k个最近邻居。

distances, indices = tree.query(points, k=1)
print(repr(indices))
# array([1, 8])

---

在这种简单的情况下，如果你的数据集和查询点集都很小，并且每个查询点都与数据集中的一行完全相同，那么使用简单的布尔运算和广播比构建和查询k-D树更快：

data, points = np.array(data), np.array(points)
indices = (data[..., None] == points.T).all(1).argmax(0)

data[..., None] == points.T将广播到一个 (nrows, ndims, npoints)数组中，对于较大的数据集而言，这可能会快速消耗内存。在这种情况下，您可以通过正常的for循环或列表推导式获得更好的性能：

indices = [(data == p).all(1).argmax() for p in points]

- ali_m