Example of searching a KD Tree with Scipy

kdtree.py

from scipy.spatial import KDTree
import numpy as np
import itertools
import matplotlib.pyplot as plt
import random


pos = np.zeros((10, 2))
for (i, j) in itertools.product(range(10), range(2)):
    pos[i,j] = random.randint(0, 10)
tree = KDTree(pos)

x = np.array([[5,5]])
(dist, indices) = tree.query(x)
print(dist, indices)
closest = np.array([pos[i,:] for i in indices])
print(closest)
plt.scatter(pos[:,0], pos[:,1])
plt.scatter(x[:,0], x[:,1], color="r")
plt.scatter(closest[:,0], closest[:,1], color="g")

plt.xlim([-1,11])
plt.ylim([-1,11])
plt.show()