How to query nearest neighbors with SciPy KDTree

Nearest-neighbor search turns a coordinate or feature collection into an index that can answer proximity questions without an explicit comparison against every row. SciPy provides KDTree for low-dimensional point sets where the matched row identity matters alongside its distance from each query point.

SciPy's KDTree accepts an n by m numeric array, with one point per row. Its query() method accepts one point or a batch whose last dimension is m, then returns distances and original-row indexes ordered from nearest to farthest.

The tree may retain the input array instead of copying it. Leave the point data unchanged after construction, or build the tree with copy_data=True when later code must modify that array.

1. Create nearest_neighbor_query.py with the required imports.

   nearest_neighbor_query.py

   import numpy as np
   from scipy.spatial import KDTree

2. Add the labeled point data below the imports.

   points = np.array(
       [
           [0.0, 0.0],
           [2.0, 3.0],
           [2.0, 4.0],
           [5.0, 4.0],
           [8.0, 2.0],
       ]
   )
   labels = np.array(["depot", "sensor-a", "sensor-b", "sensor-c", "sensor-d"])

3. Add the query coordinates below the point data.

   queries = np.array([[2.2, 3.9], [7.4, 2.2]])

4. Build the spatial index from the point rows.

   tree = KDTree(points)

5. Query the two closest indexed rows for each coordinate.

   distances, indexes = tree.query(queries, k=2)

   The default p=2 uses Euclidean distance. For two query rows with k=2, both returned arrays have shape 2 x 2.

6. Print each returned neighbor as a labeled point with its source index and distance.

   for query_number, (query, query_distances, query_indexes) in enumerate(
       zip(queries, distances, indexes), start=1
   ):
       print(f"query {query_number}: {query}")
       for rank, (distance, index) in enumerate(
           zip(query_distances, query_indexes), start=1
       ):
           print(
               f"  {rank}: {labels[index]} index={index} "
               f"point={points[index]} distance={distance:.3f}"
           )

7. Append a brute-force comparison for the first query.

   brute_force_indexes = np.argsort(
       np.linalg.norm(points - queries[0], axis=1)
   )[:2]
   assert np.array_equal(indexes[0], brute_force_indexes)
   print("\nquery 1 brute-force check: passed")

8. Run the completed nearest-neighbor query.

   $ python3 nearest_neighbor_query.py
   query 1: [2.2 3.9]
     1: sensor-b index=2 point=[2. 4.] distance=0.224
     2: sensor-a index=1 point=[2. 3.] distance=0.922
   query 2: [7.4 2.2]
     1: sensor-d index=4 point=[8. 2.] distance=0.632
     2: sensor-c index=3 point=[5. 4.] distance=3.000
   
   query 1 brute-force check: passed

   The assertion fails if KDTree.query() and the independent NumPy distance ranking select different rows for the first query.