KD-Tree Range Queries, in depth: box classification, counting with subtree sizes, the square-root bound, ball queries and updates

The query and the three node states

An orthogonal range query is an axis-aligned box: a lower and upper bound on each of the d coordinates. Three variants matter in practice. Reporting returns the matching points, and its cost must include the output size k. Counting returns how many match, and can be far cheaper than reporting. Aggregation returns a sum, minimum or other combinable value over the matches, and costs the same as counting if the tree stores the right summaries.

Every k-d tree node covers a region of space and stores a subset of points. Against a query box, a node is in exactly one of three states. Disjoint: no point below it can match, so skip it. Contained: every point below it matches, so report the whole subtree or add its count without testing a single point. Partial: some may match, so descend, or at a leaf test each point. The whole algorithm is that classification applied top-down, and its cost is governed by how many partial nodes the query creates.

Building a tree for range work

For range work, store at each node the tight bounding box of the points below it rather than relying only on the split values. A split hyperplane bounds one side; a tight box bounds all sides and shrinks around clustered data, which turns more nodes into disjoint or contained. Store the point count at each node too, and any aggregate you need.

Make leaves buckets of 8 to 64 points rather than single points. A leaf scan over contiguous coordinates is cheap and vectorises well, while every inner node costs a branch and a pointer chase. Split by index at the median of the chosen axis, so duplicate coordinates cannot produce an empty child. The code below chooses the axis with the widest spread, which adapts to elongated data; the classic analysis assumes axes are cycled in turn, which is discussed below.

class Node:
    __slots__ = ("lo", "hi", "left", "right", "idx", "count")

def build(pts, idx, leaf_size=16):
    d = len(pts[0])
    lo = tuple(min(pts[i][k] for i in idx) for k in range(d))
    hi = tuple(max(pts[i][k] for i in idx) for k in range(d))
    node = Node()
    node.lo, node.hi, node.count = lo, hi, len(idx)
    node.left = node.right = None
    if len(idx) <= leaf_size:
        node.idx = list(idx)
        return node
    axis = max(range(d), key=lambda k: hi[k] - lo[k])      # widest spread
    idx = sorted(idx, key=lambda i: (pts[i][axis], i))     # ties broken by index
    m = len(idx) // 2
    node.idx = None
    node.left = build(pts, idx[:m], leaf_size)
    node.right = build(pts, idx[m:], leaf_size)
    return node

def leaves_under(n):
    stack = [n]
    while stack:
        n = stack.pop()
        if n.idx is not None:
            yield from n.idx
        else:
            stack += [n.left, n.right]

def range_query(root, pts, qlo, qhi, count_only=False):
    """Closed box [qlo, qhi]. Returns a count, or a list of point indices."""
    d, total, out, stack = len(qlo), 0, [], [root]
    while stack:
        n = stack.pop()
        if any(n.hi[k] < qlo[k] or n.lo[k] > qhi[k] for k in range(d)):
            continue                                        # disjoint
        if all(qlo[k] <= n.lo[k] and n.hi[k] <= qhi[k] for k in range(d)):
            total += n.count                                # contained
            if not count_only:
                out.extend(leaves_under(n))
            continue
        if n.idx is not None:                               # partial leaf
            for i in n.idx:
                if all(qlo[k] <= pts[i][k] <= qhi[k] for k in range(d)):
                    total += 1
                    if not count_only:
                        out.append(i)
            continue
        stack += [n.right, n.left]                          # partial inner node
    return total if count_only else out

Sorting at every level makes this build O(n log² n); a production build uses linear-time selection (numpy.argpartition or std::nth_element) for O(n log n). The query is iterative with an explicit stack so deep trees cannot overflow the call stack.

Worked example: tracing a query

Take the ten points P0=(1,8), P1=(2,2), P2=(3,5), P3=(4,9), P4=(5,1), P5=(6,6), P6=(7,3), P7=(8,8), P8=(9,4), P9=(3,3), build with leaf_size=2, and query the closed box [2,6] x [2,6]. The trace below is what the code above actually does on this input.

The root box is (1,1) to (9,9). Spreads tie at 8, so the first axis, x, wins, and the five smallest x values {P0, P1, P2, P9, P3} go left. The left child's box (1,2) to (4,9) is taller than wide, so it splits on y into {P1, P9} and {P2, P0, P3}; the latter splits on y again into {P2} and {P0, P3}. The right side {P4..P8} splits on y into {P4, P6} and {P8, P5, P7}, and the latter into {P8} and {P5, P7}.

1. Root (1,1)-(9,9): partial. Push both children.
2. Left (1,2)-(4,9): partial. Its child (2,2)-(3,3) is contained: P1 and P9 are added without a point test.
3. Node (1,5)-(4,9): partial. Leaf {P2}, box (3,5)-(3,5), is contained: add P2. Leaf {P0, P3}, box (1,8)-(4,9), has lowest y 8 above the query's 6, so it is disjoint and skipped.
4. Right (5,1)-(9,8): partial. Leaf {P4, P6}, box (5,1)-(7,3), is partial; testing both points finds no match (P4 has y 1, P6 has x 7).
5. Node (6,4)-(9,8): partial. Leaf {P8} is disjoint (x 9 exceeds 6). Leaf {P5, P7} is partial; P5 at (6,6) matches on the closed boundary, P7 does not.

Result: P1, P9, P2, P5, a count of 4. Only four points (P4, P6, P5, P7) were tested individually, and a count-only query would have added the contained nodes' stored counts directly. On ten points the saving is small; the value is that the same rules applied to a million points open only the thin shell of nodes along the query's boundary. Note also that P5 sits exactly on the edge: decide whether your boxes are closed or half-open and make the disjoint and contained tests agree with that choice, or boundary points will be lost or double counted.

Why only a thin shell of nodes is opened

Why is the shell thin? Use the classic tree that alternates x and y splits with balanced medians and single-point leaves. Ask how many nodes a single vertical line can pass through. At a node split on x, the line lies on one side, so it enters one child. That child splits on y, and the line crosses both of its halves. So every two levels the line enters two of the four grandchildren, each holding about n/4 points: Q(n) = 2 + 2Q(n/4), which solves to O(√n).

A partial node must be crossed by one of the query box's four edges; a node that no edge crosses is either inside or outside the box. Each edge lies on a line, so at most O(√n) nodes are partial. Contained nodes are children of partial nodes, so there are O(√n) of them too, and reporting adds the output size k. That gives O(√n + k) for reporting and O(√n) for counting in two dimensions, and O(n^1-1/d + k) in d dimensions by the same argument with d levels per cycle.

Two caveats. First, the bound is a worst-case guarantee for cyclic splitting. Widest-spread splitting, as in the code above, usually does at least as well on real data, but the proof does not cover it, and adversarial inputs can make it worse. Second, n^1-1/d approaches n quickly: in 8 dimensions it is n^0.875, so for a million points the bound allows around 180,000 partial nodes. Range trees answer in O(log^d n + k) instead, at the price of O(n log^d-1 n) space; k-d trees keep linear space and accept the weaker bound.

Balls, other shapes and batches

The three-way test generalises to any region where you can bound a box. For a ball of radius r around q, a node is disjoint if the minimum distance from q to its box exceeds r, and contained if the maximum distance from q to the box (the farthest corner) is at most r.

def min_dist2(q, lo, hi):
    return sum(max(lo[k] - q[k], 0, q[k] - hi[k]) ** 2 for k in range(len(q)))

def max_dist2(q, lo, hi):
    return sum(max(abs(q[k] - lo[k]), abs(q[k] - hi[k])) ** 2 for k in range(len(q)))

# disjoint  if min_dist2(q, n.lo, n.hi) >  r * r
# contained if max_dist2(q, n.lo, n.hi) <= r * r

Half-spaces, polygons and time-window-plus-region queries work the same way: you need a cheap conservative test for disjoint and one for contained. If the contained test is hard, omitting it keeps results correct and only costs speed. This is the same pattern libraries expose as query_ball_point in SciPy's cKDTree and query_radius in scikit-learn's KDTree.

For many queries at once, such as counting neighbours of every point within r, build a second tree over the query points and traverse both together. A pair of nodes whose boxes are farther apart than r is pruned for all their points at once, and a pair entirely within r adds the product of their counts. SciPy's count_neighbors is built on this dual-tree idea.

Engineering, updates and testing

A pointer-per-node tree in Python is fine for learning and slow for production. Production trees store nodes in a flat array with children at implicit or stored offsets, store point coordinates reordered so each leaf's points are contiguous in memory, and keep per-axis coordinate arrays so a leaf scan is a few SIMD comparisons per point. Counting queries then touch only node boxes and counts, which fit in cache for trees of millions of points.

K-d trees are static. Inserting into a balanced tree breaks its balance, and deleting leaves stale boxes. Three workable policies: rebuild periodically and keep a small unsorted buffer of recent points that every query also scans; mark deletions with a tombstone and decrement counts on the path; or use the logarithmic method, keeping trees of sizes 1, 2, 4, 8 and merging equal sizes on insert, so queries visit O(log n) trees and inserts cost amortised O(log² n). If writes dominate, an R-tree is designed for them.

Always test against a brute-force oracle: random points, random boxes including empty, degenerate and fully covering ones, many duplicates, and boxes whose edges land exactly on point coordinates. Boundary handling is where k-d range code is most often wrong.

Failure modes

• Boundary inconsistency. Closed point tests with half-open box tests lose points on edges. Pick one convention and test exact-edge cases.
• Empty children from duplicates. Splitting by value at a heavily duplicated median can put every point on one side and recurse forever. Split by index.
• Skewed data defeats the bound. Highly clustered or collinear points with cyclic splits produce long thin cells that many queries cross. Tight boxes and widest-spread splits help.
• High dimensions. Beyond about 8 to 10 dimensions most nodes become partial and a vectorised scan wins. Measure nodes visited against a linear scan on your data.
• Large outputs. When k is a sizeable fraction of n, reporting is dominated by copying and any index is close to a scan; return counts or stream results instead.
• Floating point. Coordinates computed in different precisions compare unequal at boundaries; store and query in the same type.

Trade-offs and alternatives

What to do next

1. Write the build and range query above for your dimension and test them against a brute-force oracle.
2. Decide closed or half-open boxes and add exact-edge tests for that convention.
3. Store tight boxes and counts per node, and add a count-only path that never touches points.
4. Instrument nodes visited, partial nodes and points tested; compare with a linear scan on real queries.
5. Tune leaf size between 8 and 64 by measurement, and lay leaf points out contiguously.
6. If you run many queries at once, try a dual-tree traversal before scaling hardware.
7. Choose an update policy: periodic rebuild with a side buffer, tombstones, or the logarithmic method.
8. If dimensions exceed about ten or writes dominate, benchmark an R-tree or a scan before committing.