AI.Cruz - Learn AI Through Interactive Visualizations

MACHINE LEARNING

Efficient Nearest Neighbour Search: KD-Tree and Ball Tree

Brute-force KNN is O(nd) per query. KD-Trees partition space along axes; Ball Trees partition into hyperspheres. Both reduce average query time to O(log n) — the right choice depends on dimensionality.

1The Nearest Neighbour Search Problem

K-Nearest Neighbours is conceptually simple: for a query point, find the k training points with smallest distance. The naive implementation scans every training point — O(nd) per query where n is dataset size and d is dimensionality. At a million points and 100 dimensions, that is 100 million distance computations per prediction. Specialised data structures reduce this to O(log n) on average.

Brute force cost

✗ O(nd) per query — scans every point
✗ 1M points × 100 dims = 100M ops per prediction
✗ Scales poorly as dataset grows

Tree-based approaches

✓ O(log n) average query time
✓ Pre-compute spatial structure once
✓ Skip entire subtrees via triangle inequality

2KD-Tree: Axis-Aligned Space Partitioning

A KD-Tree (k-dimensional tree) is a binary tree where each node splits the data along one axis using the median value. At depth 0 it splits on axis 0 (x), at depth 1 on axis 1 (y), cycling through all dimensions. This creates a recursive partitioning of space by axis-aligned hyperplanes. Construction is O(n log n); a single approximate nearest neighbour query is O(log n).

Search algorithm: traverse the tree like a BST to reach a leaf. Store it as the current best. Unwind upward: at each node, check if the other subtree could contain a closer point by testing whether the query circle crosses the splitting plane. If it does, explore that subtree too. This backtracking is what guarantees correctness despite the greedy descent.

Construction
Sort points by the current axis, pick the median as root, recurse on left and right halves alternating axes. Result: a balanced binary tree.

Best for
Low to moderate dimensionality (d < 20). In low dimensions, most of the tree is pruned during search — very fast in practice.

Curse of dimensionality
In high dimensions (d > 20), nearly every subtree must be explored — the tree degenerates toward brute force. Use Ball Tree instead.

3Ball Tree: Hypersphere Partitioning

A Ball Tree partitions data into nested hyperspheres (balls) rather than axis-aligned boxes. Each node stores a centroid (pivot) and a radius — the distance from the pivot to the farthest point in the node. Construction: pick two far-apart points, project all data onto their connecting vector, split at the median, recurse. Ball Trees handle high-dimensional data and data on low-dimensional manifolds much better than KD-Trees.

Triangle inequality pruning: for query q and ball with centroid c and radius r, the minimum possible distance from q to any point inside the ball is max(0, dist(q,c) - r). If this minimum already exceeds the current best distance, the entire ball is skipped. This is more effective in high dimensions where axis-aligned splits fail.

Ball Tree advantages

✓ Works on low-dimensional manifolds in high-D space
✓ Effective triangle inequality pruning in high dimensions
✓ Handles non-axis-aligned cluster structure

Ball Tree limitations

✗ Slower construction than KD-Tree
✗ Sensitive to outliers (outlier stretches the ball radius)
✗ Partitions may be imbalanced

4When to Use Which

Use KD-Tree
d < 20, data is axis-aligned, you want fast construction. Typical for tabular datasets with few numeric features.

Use Ball Tree
d > 20, data lies on a manifold, or KD-Tree queries are slow. Typical for image embeddings or NLP feature spaces.

Use Brute Force
Very small datasets (n < 1000) or very high d where both trees degenerate. sklearn auto-selects based on n and d.

K-Nearest Neighbors →DBSCAN →BIRCH Algorithm →

Build a 2D KD-Tree from scratch and query nearest neighbour

import math

class KDNode:
    def __init__(self, point, left=None, right=None):
        self.point = point
        self.left = left
        self.right = right

def build_kd_tree(points, depth=0):
    if not points:
        return None
    k = len(points[0])
    axis = depth % k
    sorted_pts = sorted(points, key=lambda p: p[axis])
    mid = len(sorted_pts) // 2
    return KDNode(
        point=sorted_pts[mid],
        left=build_kd_tree(sorted_pts[:mid], depth + 1),
        right=build_kd_tree(sorted_pts[mid + 1:], depth + 1),
    )

def euclidean(a, b):
    return math.sqrt(sum((x - y) ** 2 for x, y in zip(a, b)))

def nearest_neighbour(root, query, depth=0, best=None):
    if root is None:
        return best
    dist = euclidean(root.point, query)
    if best is None or dist < best[1]:
        best = (root.point, dist)
    k = len(query)
    axis = depth % k
    diff = query[axis] - root.point[axis]
    near, far = (root.left, root.right) if diff < 0 else (root.right, root.left)
    best = nearest_neighbour(near, query, depth + 1, best)
    # Check if far subtree can contain a closer point
    if abs(diff) < best[1]:
        best = nearest_neighbour(far, query, depth + 1, best)
    return best

# Post offices in 2D (x, y)
post_offices = [
    (2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2), (10, 10)
]

tree = build_kd_tree(post_offices)
query = (6, 5)

result = nearest_neighbour(tree, query)
print("Post office locations:", post_offices)
print("Your location:        ", query)
print()
print("Nearest post office:", result[0])
print("Distance:          ", round(result[1], 3))
print()
# Verify by brute force
bf_best = min(post_offices, key=lambda p: euclidean(p, query))
print("Brute force check:  ", bf_best, "(distance " + str(round(euclidean(bf_best, query), 3)) + ")")
print("KD-Tree correct:    ", result[0] == bf_best)