Kth smallest with subtree sizes

Why this matters

Suppose a game keeps every player's score in a binary search tree (BST), and the page shows the 100th highest score. Without any extra help, the only way to find it is to list all the scores in sorted order and count down to the 100th. With a million players, that walks a million nodes on every page view.

This lesson stores one extra number in each node. The number is the size of that node's subtree, the count of nodes in it. With those counts, the program can skip whole subtrees at once. The kth smallest and kth largest queries then take time proportional to the tree height, not the number of players. The insert and delete rules from the earlier lessons in this module still apply. This lesson adds the job of keeping the counts correct when those operations run.

The idea

A subtree is a node together with everything below it. The size of a subtree is the number of nodes in it. A leaf has size 1. A node with one child whose subtree has 2 nodes has size 3, because it adds itself to the 2 nodes below.

The kth smallest key is the key with exactly k minus 1 keys below it. The first smallest key is the minimum, with 0 keys below it. The second smallest has exactly one key below it, and so on. This lesson counts from 1, the same way people count ranks in a race. The kth largest is the key with k minus 1 keys above it.

Start with a tree built by inserting 50, 30, 70, 20, 40, 60 and 80 in that order:

            50 (7)
          /        \
      30 (3)      70 (3)
     /    \       /    \
  20 (1) 40 (1) 60 (1) 80 (1)

The number in parentheses is the size of each subtree. The root has size 7 because all seven nodes are below it. The node 30 has size 3 because it has itself, 20 and 40.

Finding the 3rd smallest

Here is how the search uses the sizes. We want the 3rd smallest key, so k is 3.

1. Start at the root, 50. Its left subtree has size 3. Since k is 3, which is not more than 3, the answer is in the left subtree. Go left.
2. We are at 30. Its left subtree has size 1, the node 20. Together, that subtree and node 30 hold the 2 smallest keys, so the answer is not among them. Subtract 2 from k. The new k is 1. Go right.
3. We are at 40. Its left subtree has size 0. The k value is 1, which is one more than 0, so this node is the answer.

The answer is 40. The walk never compared the target with 20, 60, 70 or 80. It only read counts and then moved to one child.

Written as rules, the search at each node works like this. Let left be the size of the left subtree:

• If k is at most left, the answer is somewhere in the left subtree. Go left with the same k.
• If k is exactly left plus one, this node is the answer.
• Otherwise the answer is in the right subtree. The left subtree and this node account for left plus one keys, so go right with k reduced by that amount.

Finding the kth largest

The kth largest key is the same as the kth smallest counted from the other end. The kth largest is the (n minus k plus 1)th smallest, where n is the number of keys. In the tree above, n is 7. The 2nd largest is the 6th smallest. Checking the sorted list 20, 30, 40, 50, 60, 70, 80, the 6th key is 70, and the second key from the top is also 70.

Here are the first few values from the tree above:

Keeping sizes correct

The searches are only as good as the numbers stored in each node. Two operations change the sizes: insert and delete. The rule is simple. A node's size is one plus the size of its left subtree plus the size of its right subtree. Any time a subtree changes, recompute the size of every node on the path from the change back up to the root.

The code below does that recomputation in one helper named pull. Insert and erase call it on the way back up from the recursive call, so each node on the path is fixed after its child is fixed.

#include <memory>
#include <optional>
 
struct Node {
    int key;
    int size = 1;  // number of nodes in this subtree, including this node
    std::unique_ptr<Node> left;
    std::unique_ptr<Node> right;
    explicit Node(int k) : key(k) {}
};
 
int size_of(const Node* node) {
    return node ? node->size : 0;
}
 
// Recompute one node's size from its two children.
void pull(Node* node) {
    node->size = 1 + size_of(node->left.get()) + size_of(node->right.get());
}
 
std::unique_ptr<Node> insert(std::unique_ptr<Node> node, int key) {
    if (!node) return std::make_unique<Node>(key);
    if (key < node->key) {
        node->left = insert(std::move(node->left), key);
    } else if (key > node->key) {
        node->right = insert(std::move(node->right), key);
    }
    pull(node.get());
    return node;
}

Notice that insert does not add 1 to the sizes along the path. It recomputes each size from the children. The result is the same, but the recomputation is harder to get wrong. If a duplicate key arrives, nothing changes, and the recomputation gives the same numbers as before.

Delete

Deleting a node has three cases. A node with no children is removed. A node with one child is replaced by that child. A node with two children is the hard case. The usual fix is to copy the smallest key from its right subtree into the node, then delete that key from the right subtree. The key that was copied is the inorder successor, the next key in sorted order. The second lesson in this module covers these cases in more detail. Here the only new point is that every node on the path must have its size recomputed.

std::unique_ptr<Node> erase(std::unique_ptr<Node> node, int key) {
    if (!node) return nullptr;
    if (key < node->key) {
        node->left = erase(std::move(node->left), key);
    } else if (key > node->key) {
        node->right = erase(std::move(node->right), key);
    } else {
        if (!node->left) return std::move(node->right);   // zero or one child
        if (!node->right) return std::move(node->left);
        // Two children: copy the smallest key on the right, then delete that key there.
        const Node* successor = node->right.get();
        while (successor->left) successor = successor->left.get();
        node->key = successor->key;
        node->right = erase(std::move(node->right), successor->key);
    }
    pull(node.get());
    return node;
}

When a node is removed in the first two branches, the function returns its child directly, so no size is recomputed for the removed node. That is correct, because the removed node is no longer in the tree. The parent of the removed node is recomputed by its own call, once the child pointer is replaced.

Kth smallest and kth largest

With the sizes in place, the search walks one path, as described above:

// k is 1-based: k = 1 asks for the smallest key.
std::optional<int> kth_smallest(const Node* node, int k) {
    if (k < 1 || k > size_of(node)) return std::nullopt;
    while (node) {
        int left = size_of(node->left.get());
        if (k <= left) {
            node = node->left.get();              // answer is in the left subtree
        } else if (k == left + 1) {
            return node->key;                     // this node has exactly k-1 smaller keys
        } else {
            k -= left + 1;                        // skip the left subtree and this node
            node = node->right.get();
        }
    }
    return std::nullopt;
}
 
// The kth largest is the (n - k + 1)th smallest.
std::optional<int> kth_largest(const Node* node, int k) {
    int n = size_of(node);
    if (k < 1 || k > n) return std::nullopt;
    return kth_smallest(node, n - k + 1);
}

The first check handles bad input. A k of 0, a negative k, or a k bigger than the number of keys returns std::nullopt instead of walking off the tree.

Worked examples

Example 1: printing sizes and asking for ranks

This program builds the tree from the diagram. It prints each key with its size in sorted order, then asks for every rank from 1 to 7 in both directions. The helper show turns an optional into text. The helper print_sizes walks the tree in sorted order and writes each key with its size in parentheses.

#include <iostream>
#include <string>
#include <vector>
 
void print_sizes(const Node* node, std::string& out) {
    if (!node) return;
    print_sizes(node->left.get(), out);
    out += std::to_string(node->key) + "(" + std::to_string(node->size) + ") ";
    print_sizes(node->right.get(), out);
}
 
std::string show(std::optional<int> v) {
    return v ? std::to_string(*v) : std::string("none");
}
 
int main() {
    std::unique_ptr<Node> root;
    for (int k : {50, 30, 70, 20, 40, 60, 80}) {
        root = insert(std::move(root), k);
    }
 
    std::string sizes;
    print_sizes(root.get(), sizes);
    std::cout << "sizes: " << sizes << '\n';
 
    std::cout << "smallest 1..7:";
    for (int k = 1; k <= 7; ++k) std::cout << ' ' << show(kth_smallest(root.get(), k));
    std::cout << '\n';
    std::cout << "largest 1..7: ";
    for (int k = 1; k <= 7; ++k) std::cout << ' ' << show(kth_largest(root.get(), k));
    std::cout << '\n';
    std::cout << "k = 8 (out of range): " << show(kth_smallest(root.get(), 8)) << '\n';
 
    root = erase(std::move(root), 30);
    sizes.clear();
    print_sizes(root.get(), sizes);
    std::cout << "after erase 30: " << sizes << '\n';
    std::cout << "3rd smallest now: " << show(kth_smallest(root.get(), 3)) << '\n';
}

Compiled with g++ -std=c++20 -Wall -Wextra, the program prints:

sizes: 20(1) 30(3) 40(1) 50(7) 60(1) 70(3) 80(1) 
smallest 1..7: 20 30 40 50 60 70 80
largest 1..7:  80 70 60 50 40 30 20
k = 8 (out of range): none
after erase 30: 20(1) 40(2) 50(6) 60(1) 70(3) 80(1) 
3rd smallest now: 50

The largest line has two spaces after the colon, because the label ends with a space and each answer adds one more. The sizes line ends with a space, because each entry is followed by one.

Look at the erase result. Node 30 had two children, 20 and 40. Its successor is 40. The program copies 40 into the node, then removes the old 40 from the right side, which was a leaf. The node now holds 40 and has size 2, the node itself plus 20. The root drops from 7 to 6, because one key left the tree. The 3rd smallest is now 50, because the sorted list is 20, 40, 50, 60, 70, 80.

Example 2: a tree with a longer left side

Build a tree with keys 5, 2, 8, 1, 4, 7, 9 and 3, in that order. The shape is less even than the first example. The node 2 has a left child 1 and a right subtree 4 with a left child 3. Print the sizes:

sizes: 1(1) 2(4) 3(1) 4(2) 5(8) 7(1) 8(3) 9(1)

The node 2 has size 4 because it covers the keys 1, 2, 3 and 4. The root has size 8. The 6th smallest key is 7, because the sorted keys are 1, 2, 3, 4, 5, 7, 8, 9. The 2nd largest is 8. The walk for the 6th smallest starts at 5. Its left subtree has size 4, and 6 is more than 5, so skip 5 and its left subtree, which leaves k at 6 minus 5, or 1. Go right to 8. Its left subtree is the node 7, of size 1. The k value 1 is at most 1, so go left to 7. The node 7 has an empty left subtree, and k is 1, so 7 is the answer.

Now erase 5. Node 5 has two children. Its successor is 7, the smallest key on the right side. The root becomes 7, and the old 7 is removed from the right subtree. The new sizes are:

1(1) 2(4) 3(1) 4(2) 7(7) 8(2) 9(1)

The root now has size 7, which is 1 plus the left size 4 plus the right size 2. Asking for the 4th smallest gives 4, because the keys are now 1, 2, 3, 4, 7, 8, 9. A separate run of the same functions gave none for the 0th smallest, because k must be at least 1.

Example 3: the cost of a query

A search goes one step down the tree at a time. It never looks at both children of a node, so the number of steps is at most the height. The height is the number of edges on the longest path from the root to a leaf.

For a balanced tree with a million keys, the height is about 20, because 2 to the power of 20 is about a million. A kth query then takes about 20 steps. A chain of a million keys, which happens when sorted input is inserted into a plain BST, has height close to a million. The walk then takes about a million steps. The lesson "Why sorted input makes a BST a chain" explains how that shape appears. Sizes do not make the chain shorter. They only make each step decide where the answer is, without visiting the rest of the subtree.

Common mistakes

Forgetting to recompute sizes

The most common bug is to write insert without the call to pull. The code below is a copy of the program from Example 1, with the pull line removed from the insert function only:

std::unique_ptr<Node> insert_without_pull(std::unique_ptr<Node> node, int key) {
    if (!node) return std::make_unique<Node>(key);
    if (key < node->key) {
        node->left = insert_without_pull(std::move(node->left), key);
    } else if (key > node->key) {
        node->right = insert_without_pull(std::move(node->right), key);
    }
    return node;
}

After inserting the same seven keys, every size stays 1, and the root claims to have only one node. Asking for the 7th smallest key then returns none, because the function sees the tree as size 1. The output of the program with this bug is:

sizes: 20(1) 30(1) 40(1) 50(1) 60(1) 70(1) 80(1) 
smallest 1..7: 20 none none none none none none
largest 1..7:  20 none none none none none none

A check that compares the tree with a sorted std::set after every insert and erase catches this bug. With the buggy insert, the first failure comes at the second insert, because the root still reports size 1 while the set holds 2 keys.

Using 0 as the first rank

Ranks start at 1 in this lesson. If a loop starts at 0, the code asks for the minimum with k = 0 and gets none. Keep one convention for the whole program.

Getting the largest formula wrong

The kth largest is n - k + 1, not n - k. With 7 keys, the largest key is the 7th smallest, which is 7 - 1 + 1. The formula n - k would return the 6th smallest, which is 70, for the first largest query. Example 1 prints the largest keys in a row, so the error shows up as 70 where 80 should be at the start of the line.

Forgetting the sizes after a two-child delete

When a node with two children is removed, the function copies the successor key into the node and then deletes that key from the right side. The sizes along the right side change, and so does the size of the node itself. If the code only fixes the node that got the copied key and not the nodes between the root and that spot, the root's size is wrong. The pull call at the end of erase runs on every node on the path, which is why it must stay after the recursive call.

Practice

Exercise 1

Insert 5, 2, 8, 1, 4, 7, 9 and 3, in that order. Find the 6th smallest key and the 2nd largest key.

Show the answer

The sorted keys are 1, 2, 3, 4, 5, 7, 8 and 9. The 6th smallest is 7. The 2nd largest is 8. The program in Example 2 prints the same two values.

Exercise 2

On the tree from Exercise 1, erase 5. What are the new root size and the 4th smallest key?

Show the answer

Node 5 has two children. Its successor is 7, so the root's key becomes 7 and the old 7 is removed from the right side. The root's size is 1 plus the left size 4 plus the right size 2, which is 7. The sorted keys are now 1, 2, 3, 4, 7, 8 and 9. The 4th smallest is 4.

Exercise 3

Explain in your own words why the bug that skips pull gives none for the 7th smallest key, even though the tree has 7 keys.

Show the answer

The kth_smallest function starts by checking that k is not bigger than the size stored in the root. With the bug, every node stores size 1, so the root stores 1. The check sees 7 as bigger than 1 and returns none at once. The keys are still in the tree, and an inorder walk would find them. Only the stored count is wrong.

Exercise 4

Write count_smaller(root, key), which returns how many keys in the tree are strictly smaller than key. Use the sizes so that it takes time proportional to the height.

Show the answer

// Number of keys strictly smaller than key.
int count_smaller(const Node* node, int key) {
    int count = 0;
    while (node) {
        if (key <= node->key) {
            node = node->left.get();
        } else {
            count += size_of(node->left.get()) + 1;
            node = node->right.get();
        }
    }
    return count;
}

When the key is at or below the node, the node and its right subtree are not smaller, so only the left side can hold smaller keys. When the key is above the node, the node and its entire left subtree are smaller, so the function adds their size, which is the left size plus one, and then checks the right side. For the tree 50, 30, 70, 20, 40, 60 and 80, count_smaller with key 50 returns 3, and with key 65 returns 5. A test on 100 random trees of 25 keys, for every query from -1 to 100, found no wrong answers.

Exercise 5

Explain why kth_largest can ask kth_smallest for the position n - k + 1 without building a second tree that stores the keys in reverse order.

Show the answer

The sorted order is the same from either end. The kth largest is the key that has k minus 1 keys above it. In sorted order, the number of keys above a key is the total minus its position from the bottom. So the kth largest is the key whose position from the bottom is n minus k plus 1. The sizes already tell the search how many keys sit on each side, so one walk can find any position. No second tree is needed.

Summary

• Each node stores the number of nodes in its subtree. A node's size is one plus the sizes of its two children.
• The kth smallest key is found by comparing k with the left subtree size and walking one path, in time proportional to the tree height.
• The kth largest is the n - k + 1 smallest, so one search function serves both.
• Recompute sizes on the way back up after each insert and each erase. Do not add or subtract by hand, because a missed path leaves a wrong count.
• The bug that skips the recomputation is easy to make. A random check against a sorted list catches it on the first failing step.