API Reference / algorithms / neighbors / search /

ball_tree.py

Ball Tree Nearest Neighbor Search.

Classes

BallTree

class algorithms.neighbors.search.ball_tree.BallTree(NearestNeighborSearch)

Ball Tree for nearest neighbor search in high-dimensional spaces.

A Ball Tree is a binary tree where each node defines a D-dimensional hypersphere (ball) containing a subset of the data points. Compared to KD-trees, Ball Trees tend to be more efficient when the dimensionality is moderate to high because they use ball-shaped (spherical) partitions instead of axis-aligned splits.

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 Ball Tree is constructed and queried as follows:

  1. Compute the centroid and radius of the bounding ball for all
points in the current subset.
  1. If the number of points is at or below leaf_size, create a leaf
node storing the point indices.
  1. Otherwise, find the dimension with the greatest spread and split the
points along the median of that dimension.
  1. Recursively build left and right child subtrees.
  2. During a query, prune subtrees whose bounding ball cannot contain
a point closer than the current k-th nearest neighbor.

Theory

The pruning criterion for a query point q and a ball node with center c and radius r is:

d(q, c) - r \geq d_k

where d_k is the distance to the current k-th nearest neighbor. If the inequality holds, the entire subtree is skipped because no point inside the ball can be closer than d_k.

The distance metric used 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 create deeper trees with more pruning opportunities 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-dimensional or adversarial data)
  • Space: O(n) for the tree structure
When to use BallTree:
  • Moderate to high-dimensional data where KD-trees degrade
  • When the data has intrinsic low-dimensional structure (e.g., manifolds)
  • Repeated nearest-neighbor queries against a fixed dataset
  • When non-axis-aligned clusters are present in the data

References

Omohundro1989
Omohundro, S.M. (1989). Five Balltree Construction Algorithms. International Computer Science Institute Technical Report.
Liu2006
Liu, T., Moore, A.W. and Gray, A.G. (2006). New Algorithms for Efficient High-Dimensional Nonparametric Classification. Journal of Machine Learning Research, 7, pp. 1135-1158.

See Also

KDTree

Axis-aligned tree, faster in low dimensions.
BruteForceSearch

Brute-force search as a baseline.

Build a Ball Tree and query for the nearest neighbor:

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

Methods

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

Build the Ball Tree structure from training data.

Parameters
X
np.ndarray of shape (n_samples, n_features)
The training data.
Returns
self
BallTree
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.