Order-Statistic Augmentation and Rank Queries

An ordered map answers “what is the key at position k” and “how many keys are smaller than this one” only if the tree carries a count of the nodes beneath each node. The count is a single integer per node, and the reason it can be maintained is a property of rotations established earlier: a rotation changes the membership of exactly one subtree, so only the two nodes involved need their aggregates recomputed, and each recomputation is a constant-time combination of child aggregates. That is the general pattern for any augmentation, not just counts, and it is the reason augmented trees are practical rather than a theoretical curiosity.

The pattern has one hard requirement: the aggregate must be a function of the node’s own value and its children’s aggregates, and nothing else. An aggregate that depends on the order of insertion, on the position of the node in the tree, or on information not stored in the subtree cannot be maintained by local recomputation, and the whole scheme collapses.

The count and its update rule

Store at each node x a field size(x) equal to the number of nodes in the subtree rooted at x, including x itself. Then

size(x) = size(x->left) + size(x->right) + 1

with size(NULL) = 0. The value is defined by the subtree, so it is invariant under rotations that preserve the in-order sequence — a rotation changes which nodes are in which subtree, but it recomputes correctly because the formula refers only to children.

After a rotation, exactly two nodes have changed children: the node that moved down and the node that moved up. Every other node’s subtree membership is unchanged. So the update is two assignments, in bottom-up order: recompute the new child first, then the new parent. Getting the order wrong is the classic bug, because the parent’s formula reads the child’s stale value.

typedef struct node {
    int key;
    int size;                 /* nodes in this subtree, including this one */
    struct node *left, *right;
} node;

static int sz(node *t) { return t ? t->size : 0; }

/* Recompute size after a rotation; call on the lower node first. */
static void pull(node *t) { t->size = sz(t->left) + sz(t->right) + 1; }

node *rotate_right(node *x) {
    node *y = x->left;
    x->left = y->right;
    y->right = x;
    pull(x);                  /* x is now the lower node */
    pull(y);
    return y;
}

The insertion path must also recompute along the way back up, since every ancestor on the path gains a node. In an implementation that walks back to the root for balance maintenance, the recomputation is folded into the same walk. In an implementation that does not walk back — a treap, for example, which only rotates locally — the sizes along the descent path must be incremented before the rotations, since the ancestors are not revisited afterward. That difference is a real source of bugs when a tree is converted from height-balanced to randomized.

Rank and select

With counts in place, the two queries are symmetric descents.

rank(x), the number of keys strictly less than a given key, is computed by summing the sizes of the left subtrees at every step where the descent goes right, plus one for each node whose key is smaller than the target. The descent visits one node per level, so the cost is the height: O(log n) in any balanced tree, Θ(n) in a degenerate one.

select(k), the key with exactly k smaller keys, descends by comparing k against the left subtree’s size. If k < size(left), the answer is in the left subtree. If k == size(left), the answer is the current node’s key. Otherwise the answer is in the right subtree with k reduced by size(left) + 1. Again one node per level.

/* Return the key with exactly k smaller keys, 0 <= k < size(root). */
int select(node *t, int k) {
    while (t) {
        int left = sz(t->left);
        if (k < left)      t = t->left;
        else if (k == left) return t->key;
        else             { k -= left + 1; t = t->right; }
    }
    return -1;                 /* k out of range; caller must check */
}

Both queries are exact, not approximate, and both are unaffected by the balance scheme: the same code works on an AVL tree, a red-black tree, or a treap, provided the sizes are maintained. That is the practical value of augmentation — it is orthogonal to the balancing mechanism, so the choice of balance invariant is made on other grounds and the rank support is added afterward. The invariants themselves are the ones compared in red-black colouring as a height proxy and in AVL balance factors.

The count is also what makes an ordered map able to answer “the k-th element” for a k chosen at runtime, which is the operation that distinguishes an indexed sequence from a plain sorted set. A balanced tree with counts is the usual way to implement a sequence with logarithmic positional access and insertion anywhere in the middle, and it is why the structure appears in libraries that need both ordered iteration and positional access.

Interval trees and the same mechanism with a different aggregate

The interval tree is the canonical non-count augmentation. Each node stores an interval [low, high] and an aggregate max_high equal to the maximum of high over the subtree. The aggregate is recomputable from children — max_high(x) = max(x.high, max_high(x->left), max_high(x->right)) — and it is maintained through rotations by the same two-node rule.

The aggregate answers stabbing queries: given a point p, find an interval containing it. The descent compares p against the left subtree’s max_high. If p is at most that maximum, some interval in the left subtree ends at or after p, and since every interval in that subtree starts at or before the current node’s low endpoint — which the descent maintains is at most p — that interval also starts at or before p and therefore contains it. The search goes left. Otherwise the left subtree cannot contain p and the search examines the current node and continues right. The argument depends on the search order on low endpoints, which is what makes the structure an ordered tree rather than a bag.

The same pattern covers subtree sums (for prefix-sum queries), subtree minimum and maximum (for range-extremum queries), and subtree hash values (for equality checks of subtrees). Each requires that the aggregate be associative in a way that permits recomputation from the two children plus the node’s own value, and each costs one field per node.

A caution about max_high and deletion: the aggregate cannot be updated by subtracting the removed value, because a maximum does not invert. The recomputation must read the children, which is why the two-node rule is stated as a recomputation rather than as a delta. Aggregates that support inverse operations, such as sums, can be updated by delta in some cases, but the recomputation form is always correct and is what implementations should use.

Why the augmentation must be local

The requirement that the aggregate be a function of the subtree alone is not a stylistic preference; it is what makes the O(log n) bound hold. If the aggregate depended on something outside the subtree, a rotation would change it for nodes whose subtree membership did not change, and the repair would have to touch an unbounded number of nodes.

Two common proposals violate the requirement. The first is storing the depth of a node as an aggregate: depth depends on the path from the root, so a rotation changes the depth of every node in the moved subtree, and the field cannot be recomputed from children. The second is storing the count of nodes that satisfy a predicate defined relative to a query parameter, which is fine for a fixed predicate and impossible for a parameterised one — a range-count query over an arbitrary range is answered by a different technique, splitting the tree at the range boundaries and combining the pieces, not by an aggregate.

The locality requirement also bounds what a single augmentation can answer. An aggregate can answer questions about the subtree as a whole, but not questions that require combining two arbitrary subtrees in a way that depends on their relative position. Interval overlap between two separate trees, for example, is not an aggregate query.

Cost of the extra field, and when it is not worth it

The count costs one integer per node. For a red-black tree whose colour is packed into a pointer, adding a count doubles the metadata from a bit to a word, and grows the node from 24 bytes to 32 on a 64-bit machine with 8-byte keys. That is a 33 percent increase in node size, which reduces the number of nodes per cache line and increases the miss count per descent.

The trade is worth it only when the positional queries are actually used. A map that is queried by key alone does not need counts, and adding them costs memory and a recomputation on every rotation for no benefit. A map that is used as an indexed sequence — where select and rank are part of the interface — needs them, and there is no cheaper way to get O(log n) positional access in an ordered structure. The intermediate case, where rank is needed occasionally, can be served by an order-statistic tree built on demand, but rebuilding costs O(n) and the choice is usually made once at the type level rather than per query.

A subtler cost is that the count must be consistent with the tree’s contents at every moment a query can observe it. In a concurrent tree, a rotation updates two counts, and a reader that observes the tree between the two updates sees an inconsistent pair. The standard fixes are to perform the update under a lock that readers also take, or to make the counts part of an immutable node version so that a reader sees either the old or the new pair. The first is simple and serializes; the second is the technique used by persistent trees, where every modified node is a new object and the consistency problem does not arise.