AVL order statistics: maintain subtree sizes for rank and select

Last updated: 5 Oct 20269 min read

tutorial

An order-statistic tree augments each ordered node with the number of nodes in its subtree. An AVL tree also stores height to keep its shape balanced. Each insert, deletion, or rotation must refresh both fields after changing child references. Rank counts keys strictly below a requested key, even when the key is absent. Select returns the key at a zero-based ordinal by comparing that ordinal with the left subtree size and skipping the left subtree plus the current node when moving right. This implementation is a mutable set of distinct integer asset IDs; duplicate insertions and missing-key deletions raise rather than silently changing cardinality.

Operational case

Insert assets 47, 19, 61, 52, and 83. The zero-based ordinal 2 selects 52, and rank 58 is 3 because 19, 47, and 52 are smaller. Delete 47 and enumerate all remaining ordinals: 19, 52, 61, 83. A deletion can change both balance and subtree size along the path to the root. Rotating without refreshing sizes might leave the tree search ordering intact while producing the wrong rank or select result, making this a useful invariant test beyond an ordinary membership check.

Working Python program

"""AVL set with subtree sizes for zero-based select and strict rank."""

from dataclasses import dataclass


@dataclass
class OrderNode:
    asset_id: int
    left: "OrderNode | None" = None
    right: "OrderNode | None" = None
    height: int = 1
    size: int = 1


def height(node):
    return node.height if node else 0


def size(node):
    return node.size if node else 0


def refresh(node):
    node.height = 1 + max(height(node.left), height(node.right))
    node.size = 1 + size(node.left) + size(node.right)
    return node


def rotate_left(node):
    replacement = node.right
    node.right = replacement.left
    replacement.left = node
    refresh(node)
    return refresh(replacement)


def rotate_right(node):
    replacement = node.left
    node.left = replacement.right
    replacement.right = node
    refresh(node)
    return refresh(replacement)


def rebalance(node):
    refresh(node)
    difference = height(node.left) - height(node.right)
    if difference > 1:
        if height(node.left.left) < height(node.left.right):
            node.left = rotate_left(node.left)
        return rotate_right(node)
    if difference < -1:
        if height(node.right.right) < height(node.right.left):
            node.right = rotate_right(node.right)
        return rotate_left(node)
    return node


def insert(root, asset_id):
    if root is None:
        return OrderNode(asset_id)
    if asset_id == root.asset_id:
        raise ValueError("duplicate asset")
    if asset_id < root.asset_id:
        root.left = insert(root.left, asset_id)
    else:
        root.right = insert(root.right, asset_id)
    return rebalance(root)


def remove_min(root):
    if root.left is None:
        return root.right, root.asset_id
    root.left, minimum = remove_min(root.left)
    return rebalance(root), minimum


def remove(root, asset_id):
    if root is None:
        raise KeyError(asset_id)
    if asset_id < root.asset_id:
        root.left = remove(root.left, asset_id)
    elif asset_id > root.asset_id:
        root.right = remove(root.right, asset_id)
    elif root.left is None:
        return root.right
    elif root.right is None:
        return root.left
    else:
        root.right, root.asset_id = remove_min(root.right)
    return rebalance(root)


def rank(root, asset_id):
    """Number of indexed asset IDs strictly below asset_id."""
    count = 0
    while root:
        if asset_id <= root.asset_id:
            root = root.left
        else:
            count += size(root.left) + 1
            root = root.right
    return count


def select(root, ordinal):
    if ordinal < 0 or ordinal >= size(root):
        raise IndexError("ordinal outside index")
    while root:
        left_count = size(root.left)
        if ordinal < left_count:
            root = root.left
        elif ordinal == left_count:
            return root.asset_id
        else:
            ordinal -= left_count + 1
            root = root.right


asset_root = None
for asset_id in (47, 19, 61, 52, 83):
    asset_root = insert(asset_root, asset_id)
print(select(asset_root, 2), rank(asset_root, 58))
asset_root = remove(asset_root, 47)
print([select(asset_root, ordinal) for ordinal in range(size(asset_root))])

Output

52 3
[19, 52, 61, 83]

Time, space, and tradeoff

For N distinct keys, AVL height is O(log N). Insert, delete, rank, and select each follow one root-to-leaf path and take O(log N) time. Inserts and deletes use O(log N) recursion stack space; rank and select use O(1) working space. Stored nodes occupy O(N) space, with one height and one size field per node. Enumerating every ordinal by calling select repeatedly costs O(N log N), as the small trace does; a single in-order traversal would take O(N). This mutable tree is not a persistent version store and is not safe for unsynchronized concurrent writers or readers during rotation.

Common Mistakes

Do not refresh height but leave subtree size stale after rotation.
Do not confuse zero-based select with one-based rank in an API contract.
Do not decrement size when a missing-key deletion raises without changing the tree.
Do not call repeated select for every key an O(N) traversal.

Connected lessons

Use this invariant in the dispatch audit project, then check the operations quiz.

Wavelet tree: subarray counts and order statistics extends this operation contract.

data structures

range-query-structures