Find all nearest neighbors within a specific distance Find all nearest neighbors within a specific distance numpy numpy

Find all nearest neighbors within a specific distance


You could use a scipy.spatial.cKDTree:

import numpy as npimport scipy.spatial as spatialpoints = np.array([(1, 2), (3, 4), (4, 5)])point_tree = spatial.cKDTree(points)# This finds the index of all points within distance 1 of [1.5,2.5].print(point_tree.query_ball_point([1.5, 2.5], 1))# [0]# This gives the point in the KDTree which is within 1 unit of [1.5, 2.5]print(point_tree.data[point_tree.query_ball_point([1.5, 2.5], 1)])# [[1 2]]# More than one point is within 3 units of [1.5, 1.6].print(point_tree.data[point_tree.query_ball_point([1.5, 1.6], 3)])# [[1 2]#  [3 4]]

Here is an example showing how you can find all the nearest neighbors to an array of points, with one call to point_tree.query_ball_point:

import numpy as npimport scipy.spatial as spatialimport matplotlib.pyplot as pltnp.random.seed(2015)centers = [(1, 2), (3, 4), (4, 5)]points = np.concatenate([pt+np.random.random((10, 2))*0.5                          for pt in centers])point_tree = spatial.cKDTree(points)cmap = plt.get_cmap('copper')colors = cmap(np.linspace(0, 1, len(centers)))for center, group, color  in zip(centers, point_tree.query_ball_point(centers, 0.5), colors):   cluster = point_tree.data[group]   x, y = cluster[:, 0], cluster[:, 1]   plt.scatter(x, y, c=color, s=200)plt.show()

enter image description here