KD-Tree and Ball Tree: KNN Optimization

Brute-force KNN computes the distance from every query to every training sample — $O(n)$ distance computations per prediction. At one million training samples and 1000 features, that's $10^9$ multiplications per query. The KD-tree and Ball tree reduce this by organizing training points spatially so large portions of the dataset can be skipped without computing distances at all.

Anchor dataset: 7 points from the house dataset, normalized for clean spatial construction.

import numpy as np

X = np.array([
    [1.5, 2],  # A
    [2.0, 2],  # B
    [2.5, 3],  # C
    [3.5, 3],  # D
    [4.0, 4],  # E
    [4.5, 4],  # F
    [5.0, 5],  # G
])
q = np.array([3.0, 3])  # query point

The Problem: Brute Force Is O(n)

For each query, brute-force KNN:

1. Computes distance to all $n$ training samples
2. Sorts by distance: $O(n\log n)$
3. Returns the top $k$

At $n=1,000,000$ training samples and $d=1000$ features: $10^6\times 10^3=10^9$ multiplications per query. At 1ns per multiplication: 1 second per prediction. For a production API handling 1000 queries/second, this is impossible.

The key insight: if we know a query is close to region R₁, we can skip all points in region R₂ that is far from the query — without computing individual distances to points in R₂.

Building the KD-Tree

A KD-tree partitions the space with axis-aligned hyperplanes. The construction:

1. At depth $d$: split on dimension $d\mathrm{mod}p$ (cycling through dimensions)
2. Split at the median along the chosen dimension
3. Points left of median → left subtree; points right → right subtree
4. Recurse until leaf nodes have $\le$ leaf_size points

Building on 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.

The blue vertical line is the root split (D). Green horizontal lines are the left (B) and right (F) subtree splits. Six rectangular regions emerge — each leaf's bounding box.

KD-Tree Search for Nearest Neighbor of q=(3.0, 3)

Phase 1: Descend to find the leaf candidate

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.

Phase 2: 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. Check left child (A)?

   • A is in the region $x_2\le 2$. The closest point in that region to $q=(3,3)$ is at the boundary $(3,2)$: $d=\sqrt{0+1}=1.0>0.5$.
   • PRUNE: the entire left subtree of B is more than 0.5 away. Skip A.

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 (D ties C).

6. Check right subtree of D (F's subtree). Distance to the split plane ($x_1=3.5$) from $q.x_1=3.0$: $|3.0-3.5|=0.5=\text{best distance}$. Cannot prune — the boundary itself is at distance 0.5, so there might be a point exactly on it. 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 of F (E and G) since F's distance already exceeds the best.

Result: Nearest neighbor is C=(2.5, 3) (and D=(3.5, 3) ties it). 3 distance computations — instead of 7.

The Pruning Condition

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 the query to the split hyperplane is a lower bound on the distance to any point in the unexplored subtree. If even the closest possible point in that subtree (on the boundary) is farther than the current best, skip the entire subtree.

Ball Tree — Spherical Partitioning

KD-trees use axis-aligned boxes. Ball trees partition using nested hyperspheres (balls). Each node stores a center $c$ and radius $r$ such that all points in the subtree satisfy $d(x,c)\le r$.

Pruning condition: prune a ball at center $c$, radius $r$, if:

$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 minimum exceeds the best distance found so far, the entire ball is skipped.

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

Speed Benchmark

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")

~37× speedup for KD-tree, ~52× for Ball tree at $n=100,000$, $d=20$.

The Curse of Dimensionality — When KD-Tree Breaks Down

from sklearn.datasets import make_classification

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")

At $d=100$: KD-tree gives essentially zero speedup. Why? In high dimensions, all pairwise distances concentrate around the same value — the difference between the nearest and farthest neighbor vanishes. With no clear "near" vs "far" distinction, the pruning condition $|q_j-s|>d_{\text{best}}$ almost never fires; the search degenerates to visiting all nodes.

Comparison Table

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

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?