【Python】KDtree的调用

前言

查询点集中与目标点最近的k个点

from scipy.spatial import cKDTree
import numpy as np

data = np.array([[1,2],[1,3],[4,5]]) # 生成 100 个三维数据
tree = cKDTree(data) # 创建 K-D Tree
result = tree.query(np.array([5, 5]), k=2) # 查询与 [0.5, 0.5, 0.5] 最近的三个数据
print(data[result[1][0]])

参考链接

实例

#!/usr/bin/python
# -*- coding: UTF-8 -*-
import numpy as np
from matplotlib import pyplot as plt
from matplotlib.patches import Circle
from sklearn.neighbors import KDTree

np.random.seed(0)
points = np.random.random((100, 2))
tree = KDTree(points)
point = points[0]

# kNN
dists, indices = tree.query([point], k=3)
print(dists, indices)

# query radius
indices = tree.query_radius([point], r=0.2)
print(indices)

fig = plt.figure()
ax = fig.add_subplot(111, aspect='equal')
ax.add_patch(Circle(point, 0.2, color='r', fill=False))
X, Y = [p[0] for p in points], [p[1] for p in points]
plt.scatter(X, Y)
plt.scatter([point[0]], [point[1]], c='r')
plt.show()

参考链接