1. Architecture & Tight Bounding-Box Pruning Invariant
Invariant: Implicit Array Layout & Point-to-AABB Squared Distance Pruning
Instead of allocating heap nodes with left/right pointers, we reorder the input array pts[l..r] in-place around mid = (l + r) / 2 using std::nth_element along axis depth % D. Furthermore, each subtree precomputes its exact tight **Axis-Aligned Bounding Box (AABB)** $[\text{lo}_0.. \text{hi}_0] \times \dots \times [\text{lo}_{D-1}.. \text{hi}_{D-1}]$. For a query point $Q$, the minimum squared Euclidean distance from $Q$ to any point inside a subtree's AABB is:
$$d_{\min}^2(Q, \text{AABB}) = \sum_{a=0}^{D-1} \Big( \max(0LL,\; \text{lo}[a] - Q[a]) + \max(0LL,\; Q[a] - \text{hi}[a]) \Big)^2$$
During $K$-NN search with a max-heap of capacity $K$, if heap.size() == K and $d_{\min}^2(Q, \text{AABB}) \ge \text{heap.top().dist_sq}$, the entire subtree cannot improve the $K$-th nearest neighbor and is pruned immediately!
2. K-NN Branch-and-Bound Execution Flow
box_dist_sq(Q, AABB) >= heap.top() ?"} Prune -->|"Yes"| Skip["Prune Entire Subtree O(1)"] Prune -->|"No"| EvalMid["Evaluate d2 = dist_sq(Q, pts[mid])
Push to Max-Heap (Pop if size > K)"] EvalMid --> Order["Compare Q[axis] vs pts[mid][axis]"] Order -->|"Q[axis] < pts[mid][axis]"| LeftFirst["Search Left Child First,
Then Right Child (if not pruned)"] Order -->|"Q[axis] >= pts[mid][axis]"| RightFirst["Search Right Child First,
Then Left Child (if not pruned)"]
3. Why Tight AABB Pruning Beats Single-Axis Hyperplane Pruning
Full $D$-Dimensional Box Distance vs Single-Axis $\delta^2$
Standard textbook k-d trees only check the 1D perpendicular distance $\delta^2 = (Q[a] - \text{pivot}[a])^2$ to the current splitting plane. However, a deep subtree may be displaced along multiple axes $(x, y, z)$ simultaneously! Precomputing each node's tight AABB lets us sum squared clamping distances across all $D$ dimensions $\sum_{a=0}^{D-1} \delta_a^2$, pruning dead diagonal octants $2\times$–$4\times$ earlier while keeping all nodes in a flat contiguous std::vector.
4. Complexity Matrix
Trace 2D cross-product orientation tests cross(O, A, B) = (Ax-Ox)(By-Oy) - (Ay-Oy)(Bx-Ox) to prune clockwise interior vertices and detect sweep-line segment intersections.
Invariant & State
| Operation ($D$ Dimensions) | Average / Typical Time | Worst-Case Time | Space Overhead |
|---|---|---|---|
Tree Construction (std::nth_element) |
O(D * N log N) | O(D * N log N) | O(D * N) flat array |
| 1-Nearest Neighbor ($1$-NN) | O(log N) | O(N) | O(log N) stack |
| $K$-Nearest Neighbors ($K$-NN) | O(K log N) | O(N log K) | O(K) max-heap |
| Orthogonal AABB Range Search | O(N^{1 - 1/D} + M) | O(N^{1 - 1/D} + M) | O(log N) stack |
5. Complete C++20 Build Lab Implementation
#include <array>
#include <vector>
#include <cstdint>
#include <queue>
#include <algorithm>
#include <limits>
template <std::size_t D = 3>
class StaticKDTree {
public:
using Point = std::array<int64_t, D>;
struct IndexedPoint {
Point coords;
int id;
};
struct Neighbor {
int64_t dist_sq;
int id;
// Max-heap comparator: largest dist_sq sits at top() for O(1) threshold check
bool operator<(const Neighbor& o) const {
return dist_sq != o.dist_sq ? dist_sq < o.dist_sq : id < o.id;
}
};
private:
struct NodeBox {
Point lo;
Point hi;
};
std::vector<IndexedPoint> pts_;
std::vector<NodeBox> boxes_;
static int64_t point_dist_sq(const Point& a, const Point& b) {
int64_t sum = 0;
for (std::size_t i = 0; i < D; ++i) {
int64_t d = a[i] - b[i];
sum += d * d;
}
return sum;
}
static int64_t box_min_dist_sq(const Point& q, const NodeBox& box) {
int64_t sum = 0;
for (std::size_t i = 0; i < D; ++i) {
int64_t d = 0;
if (q[i] < box.lo[i]) d = box.lo[i] - q[i];
else if (q[i] > box.hi[i]) d = q[i] - box.hi[i];
sum += d * d;
}
return sum;
}
static bool box_disjoint(const NodeBox& a, const NodeBox& b) {
for (std::size_t i = 0; i < D; ++i) {
if (a.hi[i] < b.lo[i] || a.lo[i] > b.hi[i]) return true;
}
return false;
}
static bool point_in_box(const Point& p, const NodeBox& b) {
for (std::size_t i = 0; i < D; ++i) {
if (p[i] < b.lo[i] || p[i] > b.hi[i]) return false;
}
return true;
}
void build(int l, int r, int depth) {
if (l > r) return;
int mid = l + (r - l) / 2;
std::size_t axis = static_cast<std::size_t>(depth) % D;
std::nth_element(pts_.begin() + l, pts_.begin() + mid, pts_.begin() + r + 1,
[axis](const IndexedPoint& a, const IndexedPoint& b) {
return a.coords[axis] < b.coords[axis];
});
boxes_[mid] = {pts_[mid].coords, pts_[mid].coords};
for (int i = l; i <= r; ++i) {
for (std::size_t a = 0; a < D; ++a) {
boxes_[mid].lo[a] = std::min(boxes_[mid].lo[a], pts_[i].coords[a]);
boxes_[mid].hi[a] = std::max(boxes_[mid].hi[a], pts_[i].coords[a]);
}
}
build(l, mid - 1, depth + 1);
build(mid + 1, r, depth + 1);
}
void knn_rec(int l, int r, int depth, const Point& q, std::size_t k,
std::priority_queue<Neighbor>& pq) const {
if (l > r) return;
int mid = l + (r - l) / 2;
if (pq.size() == k && box_min_dist_sq(q, boxes_[mid]) >= pq.top().dist_sq) {
return; // Tight AABB branch-and-bound prune!
}
int64_t d2 = point_dist_sq(q, pts_[mid].coords);
Neighbor cand{d2, pts_[mid].id};
if (pq.size() < k) {
pq.push(cand);
} else if (cand < pq.top()) {
pq.pop();
pq.push(cand);
}
std::size_t axis = static_cast<std::size_t>(depth) % D;
bool go_left_first = q[axis] < pts_[mid].coords[axis];
if (go_left_first) {
knn_rec(l, mid - 1, depth + 1, q, k, pq);
knn_rec(mid + 1, r, depth + 1, q, k, pq);
} else {
knn_rec(mid + 1, r, depth + 1, q, k, pq);
knn_rec(l, mid - 1, depth + 1, q, k, pq);
}
}
void range_rec(int l, int r, const NodeBox& query_box, std::vector<int>& out) const {
if (l > r) return;
int mid = l + (r - l) / 2;
if (box_disjoint(boxes_[mid], query_box)) return;
if (point_in_box(pts_[mid].coords, query_box)) {
out.push_back(pts_[mid].id);
}
range_rec(l, mid - 1, query_box, out);
range_rec(mid + 1, r, query_box, out);
}
public:
explicit StaticKDTree(std::vector<IndexedPoint> pts)
: pts_(std::move(pts)), boxes_(pts_.size()) {
if (!pts_.empty()) {
build(0, static_cast<int>(pts_.size()) - 1, 0);
}
}
std::vector<Neighbor> query_knn(const Point& q, std::size_t k) const {
std::priority_queue<Neighbor> pq;
if (!pts_.empty() && k > 0) {
knn_rec(0, static_cast<int>(pts_.size()) - 1, 0, q, k, pq);
}
std::vector<Neighbor> res;
res.reserve(pq.size());
while (!pq.empty()) {
res.push_back(pq.top());
pq.pop();
}
std::reverse(res.begin(), res.end()); // Sorted closest to farthest
return res;
}
std::vector<int> query_range(const Point& lo, const Point& hi) const {
std::vector<int> out;
if (!pts_.empty()) {
range_rec(0, static_cast<int>(pts_.size()) - 1, NodeBox{lo, hi}, out);
}
return out;
}
};Production Pitfall — Using a Min-Heap Instead of a Max-Heap for $K$-Nearest Neighbors
During $K$-NN search, you want to track the $K$ smallest distances seen so far and quickly evict the largest among those $K$ candidates whenever a closer point is discovered. Therefore, the bounded priority queue of size $K$ must be a Max-Heap (std::priority_queue<Neighbor>), where pq.top().dist_sq is the current $K$-th nearest distance (the pruning radius)! Using a Min-Heap puts the 1st nearest neighbor at top(), making it impossible to prune against the $K$-th neighbor in $O(1)$ time.