Augmenting BSTs: Order-Statistic Trees (select/rank) & Interval Trees

The CLRS Augmentation Theorem & Two Canonical Augmented Trees

Why do rotations preserve local subtree summaries in $O(1)$ time? Because a rotation (rotateLeft or rotateRight) on edge $(x, y)$ only changes the descendant sets of two nodes: $x$ and $y$! The entire subtree rooted at the new top node contains the exact same set of nodes as before the rotation, so no ancestor above the rotated pair needs its associative summary recomputed from scratch (though ancestor updates along the $O(\log N)$ insertion/deletion path take $O(\log N)$ total time anyway).

1. Order-Statistic Tree (size[u] Augmentation)

Augment each node $u$ with its subtree size $\text{size}(u) = 1 + \text{size}(u.\text{left}) + \text{size}(u.\text{right})$ (where $\text{size}(\text{null}) = 0$). Then:

• select(u, k) (0-indexed $k$-th smallest key): Let $L = \text{size}(u.\text{left})$ be the number of nodes in $u$'s left subtree. If $k == L$, return $u$; if $k < L$, recurse into select(u.left, k); if $k > L$, subtract the $L + 1$ skipped elements and recurse into select(u.right, k - L - 1)!
• rank(root, key) (count of keys $< \text{key}$): Walk down from the root; every time we branch right (because $u.\text{key} < \text{key}$), add $\text{size}(u.\text{left}) + 1$ to our running accumulator.

2. Interval Tree (maxEnd[u] Augmentation for 1D Overlap Search)

Suppose we store closed intervals $I_i = [lo_i, hi_i]$ ordered primarily by their low endpoint $lo_i$ as the BST key, and augment each node $u$ with the maximum high endpoint anywhere in $u$'s subtree:

$$\text{maxEnd}(u) = \max\big(u.hi, \; \text{maxEnd}(u.\text{left}), \; \text{maxEnd}(u.\text{right})\big)$$

To find any interval in the tree that overlaps query $Q = [q_{lo}, q_{hi}]$ (recall two closed intervals overlap iff $u.lo \le q_{hi}$ AND $q_{lo} \le u.hi$), we inspect a **single root-to-leaf path** in $O(\log N)$ time:

1. If $u$ overlaps $Q$, return $u$.
2. If $u.\text{left} \ne \text{null}$ and $\text{maxEnd}(u.\text{left}) \ge q_{lo}$, go left (u = u.left).
3. Otherwise, go right (u = u.right).

Implementation: Interval Tree Overlap Search & Order-Statistic Select/Rank

// C++20: Augmented BST Primitives — Interval Overlap Search + Order-Statistic Select & Rank
#include <algorithm>
#include <cstdint>

struct AugmentedNode {
    int64_t lo, hi;      // Interval [lo, hi]; 'lo' acts as the primary BST key
    int64_t maxEnd;      // Max 'hi' in subtree
    int32_t sz = 1;      // Subtree node count for Order-Statistic rank/select
    AugmentedNode *left = nullptr, *right = nullptr;

    AugmentedNode(int64_t l, int64_t h) : lo(l), hi(h), maxEnd(h) {}
};

// O(1) Local Pull called after every child modification or rotation
inline void pullAugmented(AugmentedNode* u) noexcept {
    if (!u) return;
    u->sz = 1 + (u->left ? u->left->sz : 0) + (u->right ? u->right->sz : 0);
    int64_t m = u->hi;
    if (u->left)  m = std::max(m, u->left->maxEnd);
    if (u->right) m = std::max(m, u->right->maxEnd);
    u->maxEnd = m;
}

// 1. O(log N) Interval Overlap Search (Single root-to-leaf walk, zero backtracking!)
const AugmentedNode* findOverlappingInterval(const AugmentedNode* u, int64_t qLo, int64_t qHi) noexcept {
    while (u != nullptr) {
        if (u->lo <= qHi && qLo <= u->hi) return u; // Overlap found!
        if (u->left != nullptr && u->left->maxEnd >= qLo) {
            u = u->left;
        } else {
            u = u->right;
        }
    }
    return nullptr;
}

// 2. O(log N) Order-Statistic Select (0-indexed k-th smallest by 'lo')
const AugmentedNode* selectKth(const AugmentedNode* u, int32_t k) noexcept {
    while (u != nullptr) {
        int32_t leftSz = u->left ? u->left->sz : 0;
        if (k == leftSz) return u;
        if (k < leftSz) {
            u = u->left;
        } else {
            k -= (leftSz + 1);
            u = u->right;
        }
    }
    return nullptr;
}