KD-trees: nearest neighbors fast

Read these first if they are new to you

Nearest-neighbor search Intermediate

The idea in one paragraph

A KD-tree sorts points into boxes inside boxes. It cuts the space in two at the middle point along x, so that half the points lie on each side. Then it cuts each half along y, and so on, until every point has its own small box. To find a query's nearest neighbor, the search goes first to the box that holds the query, which gives a good candidate. Then it backs out and looks into a neighboring box only if that box could hold something closer. Most boxes fail that test, so the search skips them and every point inside them. For a hundred thousand points in 3D, a search checks a few dozen points instead of all of them.

Splitting at the middle point builds boxes inside boxes

A KD-tree, short for k-dimensional tree, is built by a rule that repeats:

Sort the points by one coordinate, starting with x.
Take the median point, the one in the middle of the sorted list. A line through it, across that coordinate, splits the points into two halves, one on each side of the line.
Repeat on each half with the next coordinate (y after x, then x again in 2D), until no points are left.

Each median point becomes a node of a tree. The first is the root; the two halves become its two branches. Each node owns a box: the region of space its branch covers.

Worked example, part 1: build a tree for seven points

The points are P1 (1, 2), P2 (3, 7), P3 (4, 4), P4 (6, 1), P5 (7, 6), P6 (8, 3) and P7 (9, 8).

Sorted by x: 1, 3, 4, 6, 7, 8, 9. The median is P4, so the root splits at x = 6.
Left half {P1, P2, P3}, sorted by y: 2, 4, 7. P3 splits it at y = 4. P1 goes below, P2 above.
Right half {P5, P6, P7}, sorted by y: 3, 6, 8. P5 splits it at y = 6. P6 goes below, P7 above.
P1, P2, P6 and P7 are each alone in a box. Each still cuts its box along x, the next coordinate, but both sides of that cut are empty, so the building stops.

The same seven points as boxes and as a tree

Each cut splits a box at its middle point. Each side gets half the points, not half the space. Blue lines split on x, green lines split on y. The label on each branch says which side of its parent's cut it holds: P5's branch holds the right half, x ≥ 6.

The search skips every box that is farther away than the best point so far

The search walks down from the root. At each node it measures the distance to that node's point and remembers the closest so far, the best. Then it goes first into the half that holds the query. When it comes back, it asks one question about the other half: could anything there be closer than the best? The other half begins at the split line. So if the query is farther from the line than from the best point, nothing in that half can win, and the search prunes it: it skips the whole box without looking at its points.

Worked example, part 2: search for the query (8, 4)

Root P4 (6, 1): squared distance 2² + 3² = 13. Best so far: P4, 13. The query's x is 8, which is at least 6, so go right first.
P5 (7, 6): squared distance 1² + 2² = 5. Best: P5, 5. The query's y is 4, below 6, so go down.
P6 (8, 3): squared distance 0² + 1² = 1. Best: P6, 1. P6 has no branches.
Back at P5: the other half lies above y = 6. The query is 2 away from that line, and 2² = 4 is more than 1. Skip P7.
Back at P4: the other half lies left of x = 6. The query is 2 away, and 4 is more than 1. Skip P1, P2 and P3.

The answer is P6, after measuring 3 of the 7 points. A check of all seven agrees: the others are at 53, 34, 16 and 17.

Which boxes does the search skip?

The tree starts as the worked example. Before you drag the query (the cross), predict which boxes turn gray when it sits in the top-left corner. Then drag it there. After that, add points and watch how the share of points checked falls.

split on x split on y skipped box nearest

Gray means never looked at. A box turns gray when the query is farther from its edge line than from the best point found so far. The dashed circle has the best distance as its radius: a gray box lies wholly beyond the line, outside it. This widget adds new points one at a time, so each new point splits the box it lands in.

kdtree.py

import numpy as np

def build(pts, depth=0):
    if len(pts) == 0:
        return None
    # a % b is the remainder of a divided by b.
    axis = depth % pts.shape[1]             # 0 splits on x, 1 on y, then x again in 2D
    pts = pts[np.argsort(pts[:, axis])]     # sort the points by that coordinate
    m = len(pts) // 2                       # the median point's index (// rounds down)
    return {"point": pts[m], "axis": axis,
            "low": build(pts[:m], depth + 1),
            "high": build(pts[m + 1:], depth + 1)}

def nearest(node, q, best=None):
    if node is None:                        # an empty branch
        return best
    p, axis = node["point"], node["axis"]
    d2 = ((p - q) ** 2).sum()
    if best is None or d2 < best[0]:        # no best point yet, or a closer one
        best = (d2, p)
    gap = q[axis] - p[axis]                 # signed distance to the split line
    # X if test else Y gives X when the test is true, and Y otherwise.
    # So near is the branch on q's side of the split line.
    near, far = (node["low"], node["high"]) if gap < 0 else (node["high"], node["low"])
    best = nearest(near, q, best)
    if gap ** 2 < best[0]:                  # could the far side hold a closer point?
        best = nearest(far, q, best)
    return best

# dtype=float stores the coordinates as decimal numbers.
pts = np.array([[1, 2], [3, 7], [4, 4], [6, 1], [7, 6], [8, 3], [9, 8]], dtype=float)
tree = build(pts)
print(tree["point"], tree["low"]["point"], tree["high"]["point"])
d2, p = nearest(tree, np.array([8.0, 4.0]))
print(d2, p)

[6. 1.] [4. 4.] [7. 6.]
1.0 [8. 3.]

A function that calls itself, as build and nearest do, is recursive: each call handles one node and hands its branches to new calls. Recursion explains how such calls run and why they stop. In practice you would use a library such as SciPy's scipy.spatial.KDTree, which is written in fast compiled code.

In 2D and 3D the search checks a few dozen points, not all of them

The table counts the distances that the code above measured per query, averaged over 200 random queries among random points spread evenly in a cube. Every answer matched brute force.

Points	Checked in 2D	Checked in 3D	Checked in 10D
1,000	16	26	751
10,000	19	33	2,965
100,000	23	37	5,508

In 2D and 3D, a hundred times more points adds only a handful of checks. The tree has about log₂ N levels, the number of times N can be halved before reaching 1: about 17 for 100,000 points. A typical search visits a few nodes per level. Building the tree costs more than a query. A good build finds each median without a full sort and takes about N × log₂ N steps: for 100,000 points, as many as 17 brute-force queries. The code above sorts again at every level, so its build costs somewhat more. So a tree pays off once you ask more than a few dozen queries of the same points.

Where KD-trees stop helping

In 10 dimensions the same search checked thousands of points, and in hundreds of dimensions it checks nearly all of them. The reason is that in many dimensions, almost every box lies close to the query along some axis. Clumped or unusual data can also make a search slower than the table suggests. Real libraries differ from this page in details: they stop splitting when a box holds a few points, they may cut at the middle of a box instead of at a point, and they test the distance to a whole box, not only to one split line, which skips a few more boxes.

Check yourself

Your turn to build: the points are (2, 6), (5, 2), (8, 5), (1, 1) and (6, 8). Which point is the root, and which points does each half get?

Answer
Sorted by x: 1, 2, 5, 6, 8. The root is (5, 2), split at x = 5. The left half gets (1, 1) and (2, 6); the right half gets (6, 8) and (8, 5). Each half is then split on y.
During a search, the best squared distance so far is 2.25 (distance 1.5). The query is 1.2 away from the split line of a branch it has not visited. Can the search skip that branch?

Answer
No. 1.2² = 1.44 is less than 2.25, so the other side of the line comes within the best distance and might hold a closer point. The search must look.
In the widget, put the query in the middle of a crowded area, then near a corner far from the points. When does the search check more points, and why?

Answer
Usually far from the points. There, the best distance stays large, so the circle crosses many split lines and fewer boxes can be skipped. Near crowded points, the best distance shrinks fast.

Go further

Optional background.

k-d tree: Nearest neighbor search Wikipedia Article Intermediate

The search in numbered steps, with an animation, and notes on cost in many dimensions.
Nearest Neighbors: K-D Tree scikit-learn user guide Article Intermediate

When a KD-tree beats brute force and when it does not, in a few paragraphs.
Benchmarking Nearest Neighbor Searches in Python Jake VanderPlas Article Advanced

Measured build and query times as the number of points and dimensions grow. From 2013, but the trends still hold.

Where this shows up

No unit or explanation builds on this one yet.