API Reference / algorithms / neighbors / search /

kd_tree.py

KD-Tree Nearest Neighbor Search.

Classes

KDTree

class algorithms.neighbors.search.kd_tree.KDTree(NearestNeighborSearch)

KD-Tree for nearest neighbor search in low-dimensional spaces.

A KD-Tree is a binary tree that recursively partitions the search space along axis-aligned splits. It is highly efficient for k-nearest neighbor queries in low to moderate dimensions (typically d < 20), but performance degrades toward brute-force speed as dimensionality increases.

This implementation delegates tree construction and querying to an optimized C++ backend with optional OpenMP parallelism, providing orders-of-magnitude speedup over pure-Python recursive traversal.

Constructor 
__init__(
    self,
    leaf_size: int = 10,
)

Overview

The KD-Tree is constructed and queried as follows:

  1. Choose the splitting dimension with the largest spread.
  2. Find the median of the data along that dimension.
  3. Partition the data: points \leq median go left, others go
right.
  1. Recursively build subtrees until the number of points in a node is
at or below leaf_size.
  1. During a query, traverse the tree starting from the closer subtree
and prune the far subtree when the splitting plane distance exceeds the current k-th nearest distance.

Theory

The pruning criterion for a query point q at an internal node splitting on dimension j with value v is:

|q_j - v|^2 \geq d_k^2

where d_k is the squared distance to the current k-th nearest neighbor. If this holds, the far subtree cannot contain a closer point and is skipped entirely.

The distance metric is Euclidean distance:

d(q, x_i) = \sqrt{\sum_{j=1}^{m} (q_j - x_{i,j})^2}

Parameters

leaf_size
int = 10
Maximum number of points in a leaf node. Smaller values lead to deeper trees and potentially faster queries but slower construction.

Attributes

X_
np.ndarray
The training data used to build the tree.
n_samples_
int
Number of training samples.
n_features_
int
Number of features in the training data.

Notes

Complexity:

  • Construction: O(n \log n) where n = number of data points
  • Query (average case): O(\log n) per query point
  • Query (worst case): O(n) per query point (high dimensions)
  • Space: O(n) for the tree structure
When to use KDTree:
  • Low-dimensional data (d < 20)
  • Repeated nearest-neighbor queries on a fixed dataset
  • When axis-aligned splits naturally separate the data
  • When construction time can be amortized over many queries

References

Bentley1975
Bentley, J.L. (1975). Multidimensional Binary Search Trees Used for Associative Searching. Communications of the ACM, 18(9), pp. 509-517. DOI: 10.1145/361002.361007
Friedman1977
Friedman, J.H., Bentley, J.L. and Finkel, R.A. (1977). An Algorithm for Finding Best Matches in Logarithmic Expected Time. ACM Transactions on Mathematical Software, 3(3), pp. 209-226. DOI: 10.1145/355744.355745

See Also

BallTree

Spherical partitions, better in higher dimensions.
BruteForceSearch

Brute-force search as a baseline.

Build a KD-Tree and query for the nearest neighbor:

python 
>>> from tuiml.algorithms.neighbors.search import KDTree
>>> import numpy as np
>>> X = np.array([[1, 2], [3, 4], [5, 6], [7, 8]])
>>> tree = KDTree(leaf_size=2)
>>> tree.build(X)
KDTree(n_samples=4, leaf_size=2)
>>> dists, indices = tree.query([3.1, 4.1], k=1)

Methods

build (self, X: np.ndarray) -> 'KDTree'

Build the KD-Tree structure from training data.

Parameters
X
np.ndarray of shape (n_samples, n_features)
The training data.
Returns
self
KDTree
Returns the instance itself.
query (self, x: np.ndarray, k: int=1) -> Tuple[np.ndarray, np.ndarray]

Find the k nearest neighbors for a query point.

Parameters
x
np.ndarray of shape (n_features,)
The query point.
k
int = 1
Number of neighbors to find.
Returns
distances
np.ndarray of shape (k,)
Distances to the k nearest neighbors.
indices
np.ndarray of shape (k,)
Indices of the k nearest neighbors in the training data.
query_batch (self, X: np.ndarray, k: int=1) -> Tuple[np.ndarray, np.ndarray]

Find the k nearest neighbors for multiple query points.

Parameters
X
np.ndarray of shape (n_queries, n_features)
The query points.
k
int = 1
Number of neighbors to find.
Returns
distances
np.ndarray of shape (n_queries, k)
Distances to the k nearest neighbors.
indices
np.ndarray of shape (n_queries, k)
Indices of the k nearest neighbors in the training data.
query_radius (self, x: np.ndarray, radius: float) -> Tuple[np.ndarray, np.ndarray]

Find all neighbors within a specified radius.

Parameters
x
np.ndarray of shape (n_features,)
The query point.
radius
float
The maximum distance to search within.
Returns
distances
np.ndarray
Distances to all neighbors within the radius.
indices
np.ndarray
Indices of neighbors within the radius.
__repr__ (self) -> str

String representation.