k-d trees for spatial search

What this chapter covers

Binary search trees organize numbers on a line so you can find one in O(log n). But a lot of data lives in space, not on a line — points on a map, pixels in color space, feature vectors in machine learning — and the questions are spatial: which stored point is nearest to this one? which points fall inside this box? A k-d tree generalizes the binary search tree to k dimensions to answer exactly those, far faster than checking every point. It splits space one axis at a time — the root divides on x, its children on y, theirs on x again — so each node carves its region in two, and a nearest-neighbor search can prune whole subtrees that are provably too far to matter. This chapter builds the tree and its pruning search, watches a query find its nearest neighbor while skipping most of the points, and confronts the structure's famous limitation: as dimensions grow, the pruning stops working, until in 16 dimensions the k-d tree examines 99% of the points — no better than brute force. That "curse of dimensionality" is one of the most important ideas in all of data science.

A bit of history

Jon Bentley introduced k-d trees in 1975, while a graduate student, as a multidimensional generalization of the binary search tree — the "k-d" is literally "k-dimensional." It was part of a wave of work on computational geometry and spatial data structures in the 1970s and 80s (quadtrees, R-trees, and others arrived around the same time) as computing began tackling geographic, graphical, and scientific data that was inherently spatial. The k-d tree became the standard structure for nearest-neighbor search in low dimensions and remains widely used (SciPy, scikit-learn, and countless graphics and robotics systems ship it). But its history is also the history of learning its limits: as researchers applied it to higher-dimensional data — especially in machine learning, where feature vectors have dozens or hundreds of dimensions — they found its efficiency collapsing, and that failure crystallized the understanding of the "curse of dimensionality," a phrase Richard Bellman had coined in 1957 for the general explosion of volume with dimension. The k-d tree is thus a double lesson: an elegant structure that works beautifully in the low dimensions it was designed for, and a concrete demonstration of why high-dimensional geometry defeats our low-dimensional intuitions.

The intuition

Building a k-d tree is like building a binary search tree, but the "key" rotates through the dimensions. At the root, split the points by their x-coordinate: pick the median-x point as the root, put points with smaller x in the left subtree and larger x in the right. At the next level, split by y instead. At the level below that, back to x. Each node thus represents a splitting plane (a vertical line in 2D at the root, then horizontal lines, alternating) that divides its region of space into two half-regions, and the tree recursively partitions the whole space into rectangular cells, one per point. Splitting at the median keeps the tree balanced, O(log n) deep.

Now the payoff, nearest-neighbor search, which is where the pruning lives. To find the point nearest a query, descend the tree toward the side of each split the query falls on — this quickly reaches the small cell the query sits in, whose point is a good first guess for the nearest. Then walk back up. At each node, you've searched the query's own side; the question is whether the other side could hold something closer. The answer is a single comparison: how far is the query from the splitting plane? If that perpendicular distance is greater than the distance to the best point found so far, then nothing on the other side of the plane can possibly be closer — the whole subtree over there is pruned, never visited. Only if the query is close enough to the plane that a closer point might lurk across it do you search the other side. In low dimensions, most subtrees get pruned this way, so the search touches O(log n) points instead of all n. The animation shows it: a query in one corner finds its nearest neighbor after visiting just 4 of 12 points, because the splitting planes let it prove the other 8 are too far to bother with.

Complexity: how it scales

Building the tree is O(n log n) (a median split at each of log n levels). Nearest-neighbor and range queries are O(log n) on average in low dimensions. But that low-dimensional caveat is the entire story of the k-d tree, and the face-off measures it directly: the fraction of all points a nearest-neighbor search actually examines, as the number of dimensions grows:

In 2 dimensions, the k-d tree examined about 0.8% of the 2,000 points to find a nearest neighbor — a hundredfold speedup over brute force, exactly the O(log n) win it's famous for. But watch the curve climb: by 8 dimensions it's examining a large fraction, and by 16 dimensions it examines 99.2% of the points — essentially all of them, no better than the brute-force O(n) scan, plus the overhead of the tree traversal. The pruning has completely stopped working. This is the curse of dimensionality made concrete, and it's not a flaw in the k-d tree — it's a fact about high-dimensional geometry that no exact spatial structure escapes. The k-d tree is superb in the low dimensions it was built for and useless in high ones, and knowing where that transition happens (roughly, once the dimension exceeds log₂ n) is essential to using it — or knowing not to.

What it's good at, what it isn't

k-d trees are the right tool for spatial queries in low dimensions — say 2 to about 10 — where their pruning gives a genuine O(log n) speedup. The classic uses: nearest-neighbor and k-nearest-neighbor search on geographic/map data (find the closest store, city, sensor), collision detection and ray tracing in graphics and games, spatial databases and geographic information systems, robotics (motion planning, point-cloud processing from LIDAR), and low-dimensional nearest-neighbor classification and regression in machine learning (scikit-learn's KNeighbors uses a k-d tree or ball tree for low-dimensional data). Range/box queries — find all points in a region — are equally efficient. When the data is genuinely low-dimensional and static-ish, the k-d tree is fast, exact, and simple.

Where it fails is high dimensions, as the face-off shows starkly: past ~10–20 dimensions it degrades to a slow brute force, so for the high-dimensional feature vectors common in machine learning (image embeddings, word vectors, hundreds of dimensions), it's the wrong structure — use approximate methods (LSH, HNSW) or reduce the dimension first. It's also best for static point sets: like a balanced BST, it doesn't rebalance on insertion, so a stream of insertions can unbalance it (rebuild periodically, or use a dynamic variant). Non-Euclidean distances and points that update frequently strain it. And for very small point sets, brute force is simpler and just as fast. The k-d tree's sweet spot is precise: low-dimensional, roughly-static spatial data, where its elegant pruning shines — and recognizing when you've left that sweet spot (the dimension count) is as important as the structure itself.

The data, or the inputs

The face-off measures the fraction of points a k-d tree examines for nearest-neighbor search as the dimension grows from 2 to 16, on 2,000 random points — the curse of dimensionality in a single curve. Correctness is checked against brute force on hundreds of random point sets across 1 to 4 dimensions: the k-d tree's nearest-neighbor must be exactly as close as brute force's (comparing distances, since ties allow different points), and its range/box search must return exactly the points a brute-force filter finds. The animation runs a 2D nearest-neighbor search: twelve points partitioned by the tree's splitting lines, and a query in the top-right corner that finds its nearest neighbor while pruning most of the tree.

Build it, one function at a time

Building the tree — median split on a cycling axis at each level:

def build_kdtree(points, depth=0):
    """Build a balanced k-d tree by, at each level, splitting the points along the current axis at
    their MEDIAN (which keeps the tree balanced, O(log n) deep). The axis cycles through the
    dimensions with depth: 0, 1, …, k-1, 0, 1, …. Left subtree holds points below the median on this
    axis, right subtree holds those above. O(n log n) to build (with a median sort per level)."""
    if not points:
        return None
    k = len(points[0])
    axis = depth % k
    points = sorted(points, key=lambda p: p[axis])
    mid = len(points) // 2
    node = KDNode(points[mid], axis)
    node.left = build_kdtree(points[:mid], depth + 1)
    node.right = build_kdtree(points[mid + 1:], depth + 1)
    return node

Nearest-neighbor search — descend, then prune subtrees that can't be closer:

def nearest(root, target, counter=None):
    """The nearest point in the tree to `target`, with pruning. Descend toward the target's side of
    each split first (it's the likely place the nearest neighbor lives); on the way back up, only
    search the FAR side of a split if the distance from the target to the splitting plane is less than
    the best distance found so far — otherwise nothing on that side can be closer, so the whole
    subtree is pruned. `counter` tallies nodes visited, to compare against brute force. Returns the
    nearest point."""
    best = [None, math.inf]                        # [point, squared distance]

    def visit(node):
        if node is None:
            return
        if counter is not None:
            counter[0] += 1
        d2 = sum((a - b) ** 2 for a, b in zip(node.point, target))
        if d2 < best[1]:
            best[0], best[1] = node.point, d2
        axis = node.axis
        diff = target[axis] - node.point[axis]     # signed distance to the splitting plane
        near, far = (node.left, node.right) if diff < 0 else (node.right, node.left)
        visit(near)                                 # explore the target's own side first
        if diff * diff < best[1]:                   # the far side could hold something closer → search it
            visit(far)                              # otherwise it's pruned entirely

    visit(root)
    return best[0]

Range search — the same pruning idea for "which points are in this box?":

def range_search(root, low, high):
    """All stored points inside the axis-aligned box [low, high] (per-dimension bounds). At each node,
    only descend into a child whose region overlaps the box — the same pruning idea applied to a
    region query. Returns the list of points in the box."""
    found = []

    def visit(node):
        if node is None:
            return
        p = node.point
        if all(low[i] <= p[i] <= high[i] for i in range(len(p))):
            found.append(p)
        axis = node.axis
        if low[axis] <= p[axis]:                    # left region might overlap the box
            visit(node.left)
        if p[axis] <= high[axis]:                   # right region might overlap the box
            visit(node.right)

    visit(root)
    return found

Watch it work

Here's a 2D k-d tree of twelve points, with the tree's recursive splits drawn as lines — vertical for x-splits, horizontal for y-splits, carving the plane into cells. The red triangle is the query, in the top-right. Watch the nearest-neighbor search: it descends toward the query's corner and quickly reaches nearby points (blue, visited; green, current best), tightening its estimate of the nearest. The key moments are the prunes: when the query is farther from a splitting line than from its best point so far, the entire region on the other side of that line is skipped — those points can't possibly be closer, so the search never even looks at them. It finds the nearest point, (9, 9), having visited only 4 of the 12 points; the splitting planes let it prove the other 8 are too far without measuring them. That proof-by-geometry is the whole power of the k-d tree — and it's exactly what the curse of dimensionality destroys in high dimensions, where the planes are always too close to prune:

The complete code

The from-scratch tab is the k-d tree — build, nearest-neighbor with pruning, and range search; the library tab is the brute-force O(n) nearest-neighbor it's checked and raced against, with a note on scipy.spatial.KDTree for production. Flip between them — brute force is trivially correct and always O(n); the k-d tree's cleverness (and its high-dimensional downfall) is entirely in the pruning test.

"""k-d tree — a binary search tree generalized to k dimensions, for organizing points in space so
that spatial queries ("which stored point is nearest to this one?", "which points lie in this
region?") run far faster than checking every point. A binary search tree splits a line of numbers;
a k-d tree splits k-dimensional space, one axis at a time, alternating axes as it descends: the root
splits on the x-coordinate, its children on y, their children on x again, and so on. Each node owns a
point and divides its region of space in two along the current axis.

The payoff is PRUNING. To find the nearest neighbor of a query point, you descend to the leaf region
the query falls in (fast, like a BST search), then walk back up — but at each node you only need to
explore the *other* side of the split if it could possibly contain something closer than the best
point found so far, which you check with a single comparison against the splitting plane. Most of the
tree gets pruned, so nearest-neighbor search is O(log n) on average in low dimensions instead of the
O(n) of checking every point. This chapter builds the tree, its nearest-neighbor search, and reveals
the catch that limits it: as dimensions grow, the pruning stops working — the "curse of dimensionality."
"""
import math


class KDNode:
    __slots__ = ("point", "axis", "left", "right")

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


# region: build
def build_kdtree(points, depth=0):
    """Build a balanced k-d tree by, at each level, splitting the points along the current axis at
    their MEDIAN (which keeps the tree balanced, O(log n) deep). The axis cycles through the
    dimensions with depth: 0, 1, …, k-1, 0, 1, …. Left subtree holds points below the median on this
    axis, right subtree holds those above. O(n log n) to build (with a median sort per level)."""
    if not points:
        return None
    k = len(points[0])
    axis = depth % k
    points = sorted(points, key=lambda p: p[axis])
    mid = len(points) // 2
    node = KDNode(points[mid], axis)
    node.left = build_kdtree(points[:mid], depth + 1)
    node.right = build_kdtree(points[mid + 1:], depth + 1)
    return node
# endregion


# region: nearest
def nearest(root, target, counter=None):
    """The nearest point in the tree to `target`, with pruning. Descend toward the target's side of
    each split first (it's the likely place the nearest neighbor lives); on the way back up, only
    search the FAR side of a split if the distance from the target to the splitting plane is less than
    the best distance found so far — otherwise nothing on that side can be closer, so the whole
    subtree is pruned. `counter` tallies nodes visited, to compare against brute force. Returns the
    nearest point."""
    best = [None, math.inf]                        # [point, squared distance]

    def visit(node):
        if node is None:
            return
        if counter is not None:
            counter[0] += 1
        d2 = sum((a - b) ** 2 for a, b in zip(node.point, target))
        if d2 < best[1]:
            best[0], best[1] = node.point, d2
        axis = node.axis
        diff = target[axis] - node.point[axis]     # signed distance to the splitting plane
        near, far = (node.left, node.right) if diff < 0 else (node.right, node.left)
        visit(near)                                 # explore the target's own side first
        if diff * diff < best[1]:                   # the far side could hold something closer → search it
            visit(far)                              # otherwise it's pruned entirely

    visit(root)
    return best[0]
# endregion


# region: range_search
def range_search(root, low, high):
    """All stored points inside the axis-aligned box [low, high] (per-dimension bounds). At each node,
    only descend into a child whose region overlaps the box — the same pruning idea applied to a
    region query. Returns the list of points in the box."""
    found = []

    def visit(node):
        if node is None:
            return
        p = node.point
        if all(low[i] <= p[i] <= high[i] for i in range(len(p))):
            found.append(p)
        axis = node.axis
        if low[axis] <= p[axis]:                    # left region might overlap the box
            visit(node.left)
        if p[axis] <= high[axis]:                   # right region might overlap the box
            visit(node.right)

    visit(root)
    return found
# endregion

"""The reference and the contrast. There's no k-d tree in the Python standard library; in production
you'd use SciPy:

    from scipy.spatial import KDTree
    tree = KDTree(points)
    dist, idx = tree.query(target)        # nearest neighbor
    idxs = tree.query_ball_point(target, r)   # all points within radius r

`brute_force_nearest` below checks every point — O(n) per query, the baseline a k-d tree beats in low
dimensions (and, tellingly, ties in high dimensions). It's both the correctness reference and the
contrast that reveals the curse of dimensionality.
"""


# region: brute
def brute_force_nearest(points, target):
    """The obvious O(n) nearest-neighbor search: compute the distance to every point and keep the
    closest. Always correct, never pruned — the reference the k-d tree is checked against, and the
    cost it improves on when the dimension is low enough for pruning to work."""
    best, best_d2 = None, float("inf")
    for p in points:
        d2 = sum((a - b) ** 2 for a, b in zip(p, target))
        if d2 < best_d2:
            best, best_d2 = p, d2
    return best
# endregion

Scratch vs library

The k-d tree teaches two things at once. First, the constructive lesson: spatial partitioning with a pruning test turns O(n) nearest-neighbor search into O(log n), by using geometry to prove most points can't be the answer — the same "eliminate what can't matter" spirit as backtracking's pruning and branch-and-bound. Second, and more profound, the cautionary lesson: that beautiful speedup evaporates in high dimensions, and the face-off's climb from 0.8% to 99% is one of the clearest demonstrations of the curse of dimensionality you can run. That curse is not a k-d-tree bug — it's a fundamental property of high-dimensional space that defeats every exact method, and it's why modern high-dimensional nearest-neighbor search (the vector databases behind semantic search and RAG, the retrieval in recommendation systems) uses approximate algorithms (HNSW, LSH) instead. Understanding the k-d tree — both its low-dimensional power and its high-dimensional collapse — is understanding why that entire industry exists. In production you'd use scipy's KDTree for low-dimensional data and a vector database (FAISS, HNSW-based) for high-dimensional; building the k-d tree yourself is what makes both the pruning and its breakdown something you've measured rather than been warned about.

Where you'll actually meet it

k-d trees run wherever low-dimensional spatial queries appear. Mapping and location services use them (and their relatives) for "nearest gas station / driver / friend" queries. Computer graphics and games use them for ray tracing, collision detection, and photon mapping. Robotics uses them for motion planning and processing LIDAR point clouds (3D). Machine learning uses them for k-nearest-neighbor classification and regression on low-dimensional features (scikit-learn's implementation), and for density estimation. Scientific computing and simulation use them for particle interactions (N-body problems, smoothed-particle hydrodynamics) and spatial statistics. Databases use spatial-index cousins (R-trees) for geographic queries. And crucially, the absence of k-d trees in high-dimensional vector search — replaced by HNSW and LSH in FAISS, Pinecone, Milvus, and every vector database — is a direct consequence of the curse this chapter demonstrates. Anywhere points in low- dimensional space must be searched, a k-d tree or a close relative is likely the tool; anywhere they're high-dimensional, its failure is why something else is.

Takeaways

A k-d tree generalizes the binary search tree to k dimensions, splitting space on a cycling axis at each level so that nearest-neighbor and range queries can prune subtrees that provably can't contain a closer point — O(log n) in low dimensions, a hundredfold speedup over brute force in 2D. But that pruning depends on the nearest neighbor being genuinely close relative to the splitting planes, which fails in high dimensions: by 16 dimensions the k-d tree examines 99% of points, no better than brute force. That collapse is the curse of dimensionality — a fundamental property of high-dimensional geometry, not a structural flaw — and it's why exact spatial methods give way to approximate ones (HNSW, LSH) for the high-dimensional data of modern machine learning.

The book's final chapter is a small, perfect gem that ties the whole advanced tier together. Reservoir sampling draws a uniform random sample from a stream of unknown, possibly infinite length — data too big to store, seen only once — using fixed memory and a single pass. It's the answer to "how do you sample fairly from something you can't hold," a question the earlier structures can't touch, and it closes the book on the tier's recurring theme: that the right idea, often a probabilistic one, makes the seemingly impossible routine.