A KD-tree (k-dimensional tree) is a binary tree that partitions k-dimensional space by alternating splitting axes. It enables fast nearest-neighbor and range queries in low dimensions — the workhorse of scipy.spatial.cKDTree and many computer-graphics systems.
| Average | Worst | Space | ||||||
|---|---|---|---|---|---|---|---|---|
| Access | Search | Insertion | Deletion | Access | Search | Insertion | Deletion | Worst |
Θ(log n) |
Θ(log n) |
Θ(log n) |
Θ(log n) |
O(n) |
O(n) |
O(n) |
O(n) |
O(n) |
At depth d, split by axis d mod k. Build: pick the median along that axis as the root, recurse on each half. Nearest-neighbor query: descend to the leaf containing the target; on the way back up, check whether the bounding hyperplane could contain a closer point and recurse into the other side if so. O(log n) average for low k; degrades sharply when k > ~20 (curse of dimensionality).
class Node:
__slots__ = ("point", "axis", "left", "right")
def __init__(self, point, axis):
self.point, self.axis = point, axis
self.left = self.right = None
def build(points, depth=0):
if not points: return None
k = len(points[0])
axis = depth % k
points.sort(key=lambda p: p[axis])
mid = len(points) // 2
n = Node(points[mid], axis)
n.left = build(points[:mid], depth + 1)
n.right = build(points[mid + 1:], depth + 1)
return n
def _dist2(a, b):
return sum((x - y) ** 2 for x, y in zip(a, b))
def nearest(node, target, best=None):
if node is None: return best
d = _dist2(node.point, target)
if best is None or d < best[1]:
best = (node.point, d)
diff = target[node.axis] - node.point[node.axis]
near, far = (node.left, node.right) if diff < 0 else (node.right, node.left)
best = nearest(near, target, best)
if diff * diff < best[1]:
best = nearest(far, target, best)
return best
pts = [(2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2)]
tree = build(pts)
print(nearest(tree, (9, 2))) # ((8, 1), 2)faiss).Building the tree for the points (2, 3), (5, 4), (9, 6), (4, 7), (8, 1), (7, 2) starts by sorting on x and taking the median, (7, 2), as the root. Points with smaller x go left, larger x go right. At depth 1 each half is sorted on y: the left half yields (5, 4) with children (2, 3) and (4, 7), the right half yields (9, 6) with child (8, 1). For the query (9, 2), the search compares x = 9 with the root's 7 and goes right, compares y = 2 with 6 and goes left, and reaches (8, 1) at squared distance 2. Unwinding, it checks whether the splitting planes are closer than the best distance found: the plane y = 6 is 16 units squared away and the plane x = 7 is 4 units squared away, both farther than 2, so the entire left subtree and the other branches are pruned.
Each node represents an axis-aligned region of space. The search first walks to the leaf region that contains the query, because the nearest point is likely nearby, which quickly gives a tight best distance. While unwinding the recursion, it compares the squared distance from the query to the splitting hyperplane (diff * diff) with the best squared distance. If the plane is farther than the best point, no point on the other side can be closer and that subtree is skipped. With well-spread data in low dimensions this prunes most of the tree and visits O(log n) nodes on average.
The same structure supports range search (report all points inside a box) and k-nearest-neighbour search by keeping a bounded max-heap of the k best candidates instead of a single best.
import numpy as np
from scipy.spatial import cKDTree
rng = np.random.default_rng(0)
pts = rng.random((100_000, 3)) # 100k points in 3-D
tree = cKDTree(pts, leafsize=16)
dist, idx = tree.query([0.5, 0.5, 0.5], k=5) # 5 nearest
near = tree.query_ball_point([0.5, 0.5, 0.5], r=0.02) # radius search
pairs = tree.query_pairs(r=0.005) # all close pairs
print(idx, len(near), len(pairs))Production implementations store several points per leaf (leafsize) and use iterative traversal and vectorised distance computations; a hand-written recursive tree in Python is useful for learning but far slower. scikit-learn's KNeighborsClassifier chooses between brute force, KD-tree and ball tree automatically with algorithm="auto".
| Method | Good for | Dimensions | Exact? |
|---|---|---|---|
| KD-tree | NN and range queries on static points | low (roughly up to 10 to 20) | yes |
| Ball tree | non-Euclidean metrics, clustered data | moderate | yes |
| Brute force (vectorised) | small n or very high d | any | yes |
| HNSW / IVF (faiss, vector DBs) | embeddings, millions of vectors | high (hundreds+) | approximate |
| R-tree | rectangles and geographic shapes | 2 to 3 | yes |
KD-trees are ideal for geometric problems in 2 or 3 dimensions: GIS points, collision detection, point clouds, colour quantisation. For text or image embeddings with hundreds of dimensions, the pruning test almost never succeeds, so the search visits nearly every node; approximate nearest-neighbour indexes are the right tool there.