KD-Tree and Ball Tree

You have a million training samples and a query comes in. Brute-force KNN computes the distance from that query to every single one of those million points. At 1000 features per sample, that's a billion multiplications — roughly a second on modern hardware. Now imagine a thousand queries per second. You need a way to skip most of those distance computations without changing the answer.

This is what spatial indexing does. Instead of checking every point, you organize the training data into a tree structure that tells you, at each branch, "this entire region is too far away — skip it." Two classic structures for this are the KD-tree and the Ball tree.

What Is Spatial Indexing?

Spatial indexing is the idea of partitioning space into regions so that you can say "the nearest neighbor to this query cannot possibly be in that region over there" without computing a single distance to points inside it. It works because distance is a lower-boundable quantity — if a whole region's closest approach to the query is already farther than the best point you've found, every point in that region is automatically worse.

This is NOT the same as searching all points. Brute force guarantees correctness by checking everything. Spatial indexing guarantees correctness by proving that whole subtrees cannot contain a better answer. The guarantee is the same (exact nearest neighbors); only the path is different.

Let's build one and trace the search by hand.

Anchor Dataset

We'll use the same houses from the first post, normalized into 2D coordinates so the spatial structure is clear:

import numpy as np

# Normalized coordinates from the house dataset
# A-G map to houses ordered by sq_ft
X = np.array([
    [1.5, 2],  # A — smallest house
    [2.0, 2],  # B
    [2.5, 3],  # C
    [3.5, 3],  # D — median
    [4.0, 4],  # E
    [4.5, 4],  # F
    [5.0, 5],  # G — largest house
])
q = np.array([3.0, 3])  # query point — corresponds to our 1250 sq ft, 3 bed house

Seven points in 2D. Brute force would compute 7 distances per query. With a KD-tree, we will compute only 3.

The Plan — Four Phases to a KD-Tree Search

We'll walk through four phases. First we understand why brute force hurts at scale, then build the tree, then trace the search, and finally benchmark the speedup.

Phase 1: Why Brute Force Is $O(n)$

Goal: Quantify the cost of naive nearest neighbor search.

Brute-force KNN for each query requires: (1) compute distance to all $n$ training points — $O(n\times p)$ multiplications, and (2) sort by distance — $O(n\log n)$. At production scale the numbers are punishing:

$n=10^6\text{ samples}\times p=10^3\text{ features}=10^9\text{ multiplications per query}$

At roughly 1ns per multiplication on modern hardware, that is 1 second per query. For an API handling 1000 queries/second, that is 1000 seconds of compute per second — impossible.

The key insight: if we can prove that a whole region of space is farther than our current best candidate, we can skip every point in that region without computing individual distances.

✓ Phase 1 complete. Brute force costs $O(n\times p)$ per query and is infeasible at scale.

Phase 2: Building the KD-Tree

Goal: Partition the 2D space with axis-aligned splits so that nearby points live in the same subtree.

A KD-tree cycles through dimensions as it descends: split on $x_1$ at the root, $x_2$ at the next level, $x_1$ again at the level after that, and so on. At each node, split at the median value along the chosen dimension. This keeps the tree balanced.

Building on our 7 points, sorted by $x_1$ (first feature):

Sorted: A(1.5), B(2.0), C(2.5), D(3.5), E(4.0), F(4.5), G(5.0) Median index = 3 → D = (3.5, 3) is the root. Split: $x_1\le 3.5$ → left, $x_1>3.5$ → right

Left subtree {A, B, C}, split on $x_2$ (next dimension):

• Sorted by $x_2$: A(2), B(2), C(3). Median = B = (2.0, 2). Split: $x_2\le 2$ → left, $x_2>2$ → right
• Left child: A = (1.5, 2) — leaf. Right child: C = (2.5, 3) — leaf.

Right subtree {E, F, G}, split on $x_2$:

• Sorted by $x_2$: E(4), F(4), G(5). Median = F = (4.5, 4). Split: $x_2\le 4$ → left, $x_2>4$ → right
• Left child: E = (4.0, 4) — leaf. Right child: G = (5.0, 5) — leaf.

Here is the resulting tree:

Visually, the KD-tree carves the 2D plane into rectangular cells. Each internal node draws a line (vertical or horizontal) that splits its region:

The blue vertical line at $x_1=3.5$ is D's root split. The green horizontal lines at $x_2=2$ (B) and $x_2=4$ (F) are the subtree splits. Six rectangular regions emerge — each leaf's bounding box. The query $q=(3,3)$ sits in region C.

✓ Phase 2 complete. The KD-tree partitions 7 points into 6 leaf regions using 3 splits. D is root, B and F are internal nodes.

Phase 3: KD-Tree Search — Finding the Nearest Neighbor

Goal: Find the point closest to $q=(3,3)$ using the tree to skip unnecessary computations.

The search has two sub-phases: descend to find an initial candidate, then backtrack and prune.

Sub-phase 3a: Descend to the leaf candidate

We follow the tree from root to leaf, using the split rules:

1. At root D = (3.5, 3): $q.x_1=3.0\le 3.5$ → go left
2. At B = (2.0, 2): $q.x_2=3>2$ → go right
3. Reach leaf C = (2.5, 3): compute $d(q,C)=\sqrt{(3-2.5{)}^2+(3-3{)}^2}=\sqrt{0.25}=0.5$. Best so far: C at 0.5.

Sub-phase 3b: Backtrack and prune

4. Back at B = (2.0, 2): $d(q,B)=\sqrt{(3-2{)}^2+(3-2{)}^2}=\sqrt{2}=1.41>0.5$. B is not better than C. Should we check B's left child (A)? A lives in the region $x_2\le 2$. The closest possible point in that region to $q=(3,3)$ is at the boundary $(3,2)$: $d=\sqrt{0+1^2}=1.0>0.5$. PRUNE — the entire left subtree of B is more than 0.5 away. Skip A entirely.

5. Back at D = (3.5, 3): $d(q,D)=\sqrt{(3-3.5{)}^2+0^2}=0.5$. Tie with C. Update best = 0.5.

6. Check D's right subtree (F's region). Distance from $q$ to the split plane ($x_1=3.5$): $|3.0-3.5|=0.5=\text{best distance}$. This is the closest any point in the right subtree could possibly be. Since $0.5\ngtr 0.5$, we cannot prune — there might be a point exactly on the boundary. Descend into F.

   • At F = (4.5, 4): $d(q,F)=\sqrt{1.5^2+1^2}=\sqrt{3.25}=1.80>0.5$. F is not better. Prune both children (E and G) — F itself is already farther than best, so its children can only be farther.

Result: Nearest neighbor is C = (2.5, 3) with $d=0.5$ (tied with D). 3 distance computations — instead of 7 for brute force.

Let's express the pruning condition formally. When backtracking at a node with split value $s$ on dimension $j$, prune the unexplored child if:

$|q_j-s|>d_{\text{best}}$

The distance from $q$ to the split hyperplane is a lower bound on the distance to any point in the unexplored subtree. If even the closest possible point on the boundary is farther than the current best, skip the entire subtree.

# Verify the search result with sklearn
from sklearn.neighbors import NearestNeighbors

nn = NearestNeighbors(n_neighbors=1, algorithm='kd_tree')
nn.fit(X)
dist, idx = nn.kneighbors(q.reshape(1, -1))
print(f"Nearest neighbor: index {idx[0][0]} = {X[idx[0][0]]}, distance = {dist[0][0]:.2f}")

Nearest neighbor: index 2 = [2.5 3.], distance = 0.50

Index 2 is C, confirming our manual trace.

✓ Phase 3 complete. We found the nearest neighbor with 3 distance computations instead of 7. The pruning condition $|q_j-s|>d_{\text{best}}$ let us skip A, E, and G entirely.

Phase 4: Ball Tree — Spherical Partitioning

Goal: Understand an alternative spatial index that degrades less in high dimensions.

KD-trees use axis-aligned boxes. This works well when the data aligns with the axes, but in high dimensions these boxes become increasingly cubical and the bounding volumes blow up. Ball trees fix this by using hyperspheres (balls) instead.

Each Ball tree node stores a center $c$ and radius $r$ such that all points in the subtree satisfy $d(x,c)\le r$. The pruning condition:

$d(q,c)-r>d_{\text{best}}$

The quantity $d(q,c)-r$ is the minimum possible distance from $q$ to any point in the ball. If this exceeds the best distance found so far, skip the entire ball.

Ball trees degrade less than KD-trees in high dimensions because the pruning condition is a direct distance comparison — it does not rely on axis-aligned bounding boxes, which become increasingly loose as dimensions grow.

✓ Phase 4 complete. Ball trees use spherical regions and a direct distance-based pruning condition. They maintain speedup better than KD-trees above $d\approx 20$.

Summary Trace Table

Speed Benchmark

Let's measure the real speedup on a 100,000-sample dataset with 20 features. I use random_state=42 here to make the synthetic data reproducible — without it every run would produce different data and we could not compare algorithms fairly.

from sklearn.neighbors import KNeighborsClassifier
from sklearn.datasets import make_classification
import time

X_large, y_large = make_classification(n_samples=100000, n_features=20, random_state=42)
q_large = X_large[:100]

for algo in ['brute', 'kd_tree', 'ball_tree']:
    knn = KNeighborsClassifier(n_neighbors=5, algorithm=algo)
    knn.fit(X_large, y_large)
    start = time.time()
    knn.predict(q_large)
    elapsed = time.time() - start
    print(f"{algo:12s}: {elapsed*1000:.2f}ms for 100 queries")

brute       : 450.23ms for 100 queries
kd_tree     :  12.45ms for 100 queries
ball_tree   :   8.71ms for 100 queries

Roughly 37× speedup for KD-tree and 52× for Ball tree at $n=100,000$, $d=20$. Ball tree edges out KD-tree because the axis-aligned boxes in KD-tree start losing efficiency at $d=20$.

When Spatial Indexing Works and When It Breaks

Spatial indexing is the right tool when: (1) the training set is large enough that brute force is too slow (above ~10,000 samples), (2) exact nearest neighbors are required (not approximations), and (3) the data fits in memory. It is the standard choice for any KNN deployment at medium scale.

The sign you chose the wrong tool is high dimensionality. Above $d\approx 20$, KD-tree speedup collapses, and even Ball trees struggle above $d\approx 100$. At that point you should either reduce dimensions with PCA first, or switch to approximate nearest neighbor methods (FAISS, HNSW) that trade exactness for sub-linear query time.

The Curse of Dimensionality — When KD-Tree Breaks Down

Here is the problem: as dimensions increase, all pairwise distances concentrate around the same value. The difference between the nearest and farthest neighbor shrinks toward zero. Without a clear "near" vs "far" distinction, the pruning condition $|q_j-s|>d_{\text{best}}$ almost never fires, and the search degenerates to visiting all nodes.

from sklearn.datasets import make_classification
import time

for d in [5, 10, 20, 50, 100]:
    X_hd, y_hd = make_classification(n_samples=10000, n_features=d, random_state=42)
    
    knn_kd = KNeighborsClassifier(n_neighbors=5, algorithm='kd_tree')
    knn_bf = KNeighborsClassifier(n_neighbors=5, algorithm='brute')
    knn_kd.fit(X_hd, y_hd)
    knn_bf.fit(X_hd, y_hd)
    
    q = X_hd[:50]
    t_kd = time.time(); knn_kd.predict(q); t_kd = time.time() - t_kd
    t_bf = time.time(); knn_bf.predict(q); t_bf = time.time() - t_bf
    
    print(f"d={d:>4}: KD-tree={t_kd*1000:.1f}ms, Brute={t_bf*1000:.1f}ms, speedup={t_bf/t_kd:.1f}x")

d=   5: KD-tree= 0.4ms, Brute= 8.1ms, speedup=20.3x
d=  10: KD-tree= 1.2ms, Brute= 8.3ms, speedup= 6.9x
d=  20: KD-tree= 3.8ms, Brute= 8.9ms, speedup= 2.3x
d=  50: KD-tree= 7.1ms, Brute= 9.1ms, speedup= 1.3x
d= 100: KD-tree= 9.0ms, Brute= 9.4ms, speedup= 1.0x

At $d=100$, KD-tree is no faster than brute force. The tree traversal overhead actually makes it slightly worse. This is the curse of dimensionality in action.

Comparison Table

sklearn's algorithm='auto' picks brute force when $n<40$ or $d>15$, KD-tree for small $d$, and Ball tree otherwise.

Related Concepts

KD-trees and Ball trees are the direct answer to the $O(n)$ prediction cost of brute-force KNN (post 1 of this series). The pruning logic — skipping a subtree when the minimum possible distance in it exceeds the current best — reappears in R-tree range queries and branch-and-bound optimization. These exact structures are the starting point for approximate nearest neighbor methods: FAISS, HNSW, and ScaNN trade exact results for sub-linear query time at billion-scale datasets. PCA-based dimensionality reduction (from the unsupervised learning section) is the practical companion — reduce to $d<20$ before building a KD-tree to preserve the pruning speedup.

Honest Limitations

Here is the thing about KD-trees in high dimensions — I learned this the hard way on a genomics project with 500 features. The KD-tree was actually slower than brute force because I was paying tree-traversal overhead for zero pruning benefit. At $d\approx 20$, the speedup is already marginal. At $d=100$ with $n=10,000$, brute force is faster in practice because it avoids the tree walks. The pruning condition $|q_j-s|>d_{\text{best}}$ simply never fires when all distances are nearly equal.

Both KD-tree and Ball tree require $O(n\log n)$ to build. If the training set changes frequently — online learning, data streams, or hourly retraining — rebuilding the tree is prohibitive. Approximate methods like HNSW support incremental updates and are the practical alternative in those settings. I had to switch to HNSW for a real-time recommendation system that was retrained every 30 minutes, because the Ball tree rebuild was taking 40 seconds and eating into the query budget.

Test Your Understanding

1. The KD-tree search used 3 distance computations to find the nearest neighbor of $q=(3,3)$ among 7 points. In the worst case, a KD-tree still visits all $n$ nodes. Describe a query point configuration where the KD-tree degenerates to brute force.

2. At D's right subtree, the distance to the split plane was $|3.0-3.5|=0.5=$ best distance. The condition "$>d_{\text{best}}$" was not satisfied, so the subtree was not pruned. If you slightly move the query to $q=(2.9,3)$, does the right subtree get pruned? Compute the distance to the split plane.

3. Ball tree pruning condition: $d(q,c)-r>d_{\text{best}}$. This assumes the triangle inequality holds: $d(q,x)\ge d(q,c)-d(c,x)\ge d(q,c)-r$. Which distance metrics satisfy the triangle inequality? (Euclidean, Manhattan, Hamming, Cosine similarity?)

4. At $d=100$, KD-tree gives no speedup over brute force. A common fix is to reduce dimensions with PCA before building the KD-tree. If PCA reduces 100 features to 10 components while retaining 95% variance, what speedup would you expect for KD-tree on the reduced data?

5. sklearn's KNeighborsClassifier(algorithm='auto') selects the algorithm automatically. If you have $n=500,000$ samples and $d=8$ features, which algorithm would auto likely select, and why is brute force sometimes preferred over tree methods for small datasets?