Problem: Order Statistics

A TreeMap keeps its keys in order, so it makes sense to ask for a key by its position in that order. Which username comes fifth in alphabetical order? How many scores on a leaderboard are below this one? These questions are about a key’s rank: how many keys come before it.

The two queries are collected in an interface that extends Map:

public interface RankedMap<K extends Comparable<K>, V> extends Map<K, V> {
  K select(int k);  // the k-th smallest key, counting from 1
  int rank(K key);  // the number of keys smaller than key
}

select counts from $1$ , so select(1) is the smallest key and select(size()) is the largest. It throws IllegalArgumentException when k is less than $1$ or greater than size(). rank takes any key, stored or not, and throws IllegalArgumentException for a null key.

Take this tree. The picture shows only the keys:

             ╭────╮
             │ 50 │
             ╰──┬─╯
       ┌────────┴────────┐
    ╭──┴─╮            ╭──┴─╮
    │ 30 │            │ 70 │
    ╰──┬─╯            ╰──┬─╯
   ┌───┴───┐       ┌─────┴─────┐
╭──┴─╮  ╭──┴─╮  ╭──┴─╮      ╭──┴─╮
│ 20 │  │ 40 │  │ 60 │      │ 80 │
╰────╯  ╰────╯  ╰────╯      ╰──┬─╯
                           ┌───┴───┐
                        ╭──┴─╮  ╭──┴─╮
                        │ 75 │  │ 90 │
                        ╰────╯  ╰────╯

The keys in order are 20, 30, 40, 50, 60, 70, 75, 80, 90. So select(1) is 20 and select(6) is 70. rank(50) is $3$ , because three keys are smaller than 50. rank(72) is $6$ : 72 is not stored, and six keys are smaller than it. rank(5) is $0$ .

The obvious approach uses the map’s iterator, which visits the keys in ascending order. For select(k), take $k$ steps and return the key you reach. This is how the search tree chapter’s practice found the k-th smallest value: an in-order traversal that stops at the k-th visit. For rank(key), count steps until you reach a key that is not smaller than key. Both are correct, but both take a step per key they pass, so select(n) and the rank of a large key cost $O (n)$ . We want $O (lo g n)$ for every k and every key.

You may add fields to Node, as long as put and remove still run in $O (lo g n)$ .

Try it before reading the solution.