An implementation notebook · KNN / KD-tree
Inside a KD-tree.
Move one decision at a time. See how rows become groups, how groups become bounds, and how a heap turns those bounds into permission to skip work.
Find one position without sorting every position.
A trial pivot earns its rank by comparisons. Watch the row indices move while the original feature values stay attached to their rows.
Try Back after a swap. The row ID travels with its value. No row is deleted or replaced.
The rule that makes this possible
A count gives the pivot its rank.
If exactly three rows precede a pivot, the pivot belongs in the fourth position. Those three rows can remain internally unordered.
For equal values, compare original row IDs next. This reproduces the ordering rule in sklearn’s partition comparator.
What does sklearn actually request?
std::nth_element(
start_of_indices,
required_position,
end_of_indices,
compare_rows
);std names the C++ standard library’s namespace. :: accesses a name inside it. The function rearranges indices in place.
Changing a setting restarts the traces. Query coordinates affect search; the fitted tree depends only on the training points and leaf size.
One dataset. A permutation. Arrays of nodes.
Step through construction. Select a row in the node table to inspect the corresponding group and its tight bounding box.
The four stored arrays
| Original row | data: x₁, x₂ | Class |
|---|
data keeps training-row order. The index permutation chooses which rows belong to a node.
idx_array
node_data + node_bounds
| Node | Index slice | Leaf? | Lower → upper |
|---|
Children of node i are 2i + 1 and 2i + 2. The node record has no stored split axis or threshold; query traversal uses the child boxes.
Construction details to connect to the Markdown
Sklearn preallocates a balanced tree depth from the point count and leaf size. Ordinary leaves contain between leaf_size and twice leaf_size points, with an exception when the whole dataset is smaller. The split feature has the largest observed range; ties select the first feature. The row comparator breaks equal-coordinate ties by original row index.
This simulation uses a transparent Quickselect to produce a valid partition. The exact order within a leaf can differ from a compiled library’s result. Sklearn also scans coordinates separately when constructing bounds and when finding the split feature; the view shows the resulting geometry, not individual build-loop instructions.
Prove when an entire group can be skipped.
Every bound and heap value below is a squared distance. Click the plot to move the query, or change its coordinates above.
The fixed-size max heap
The root is the farthest retained candidate. Infinity represents an unfilled slot, so the threshold stays infinite until k candidates exist.
| Heap slot | Training row | Squared distance |
|---|
Tree traversal
Returning from one branch resumes the pending sibling. The query stays at the same coordinates.
Ties, zero distances, and the output boundary
The candidate heap rejects a new squared distance greater than or equal to its root. Node pruning uses a strict greater-than comparison. A box bound equal to the threshold can therefore be explored even when tied point candidates are rejected. Try the distance-tie preset with k = 1 and compare leaf sizes 1 and 4.
Changing tree layout can change which equal-distance rows are retained. The final order among ties is not a promise to reproduce sklearn’s sorting routine. For inverse-distance voting, selected exact matches receive weight 1 and selected nonzero-distance points receive 0. Otherwise weights use 1/distance, not 1/squared-distance.
In the audited source, default output sorting and conversion from squared to true distances happen even when only indices are returned. This view uses explicit query coordinates; the estimator’s no-argument self-exclusion convention is explained in the companion note.
Avoided work must pay for the index.
Separate the number of operations from their execution cost. Point counts explain pruning; they are not elapsed-time benchmarks.
Leaf size on the current point set
| leaf_size | Nodes | Point distances | Box bounds |
|---|
A smaller leaf is useful when extra bounds eliminate enough point checks. A deeper tree can perform the same point work with more traversal.
When geometry stops helping
Recorded in installed sklearn 1.8.0: 2,048 uniform training points, 64 queries, leaf_size = 16, seed 42 reset for each dimension.
The twenty-dimensional examples checked every training point. These observations illustrate particular datasets; they are not a universal dimensionality threshold or a timing claim.
What dense Euclidean “auto” chooses
This is the audited heuristic for ordinary dense features and unweighted Euclidean distance. The estimator does not time several backends. Sparse input, precomputed distances, and other metrics add separate rules.
When the input needs another buffer
Model of direct public KDTree storage: float64 coordinates, 8-byte indices, and the observed 32-byte node record. Counts array payload while retaining the input; excludes Python objects, labels, query heaps, and allocation overhead. A C-contiguous float64 input can share the tree buffer.
Connect these views to the remaining implementation details
Compiled feature loops reduce the price of a point check. Query batches can be split across threads; one query’s branches are not split by that estimator path. Direct KDTree queries also offer a priority queue of nodes and a second tree over queries. The companion note explains these options, optimized brute force, dtype behavior, and prediction conventions.