Solution: select

At a node, the size of the left subtree tells us how many keys in this subtree are smaller than the node’s key. Call it leftSize. Then the node’s key is the (leftSize + 1)-th smallest key in its subtree. Compare k with that:

  • If k is at most leftSize, the k-th smallest key is in the left subtree, and it is the k-th smallest there too.
  • If k is leftSize + 1, it is this node’s key.
  • Otherwise it is in the right subtree. When we go right, we skip the leftSize keys of the left subtree and the node’s own key, so we look for the (k - leftSize - 1)-th smallest key there.
@Override
public K select(int k) {
  if (k < 1 || k > size()) {
    throw new IllegalArgumentException();
  }
  return select(root, k).key;
}

// Assumes k is from 1 to size(node). Returns the node with the k-th smallest
// key in the subtree rooted at node.
private Node<K, V> select(Node<K, V> node, int k) {
  int leftSize = size(node.left);
  if (k <= leftSize) {
    return select(node.left, k);
  } else if (k == leftSize + 1) {
    return node;
  } else {
    return select(node.right, k - leftSize - 1);
  }
}

The public method checks k once, so the helper can assume it is in range. A k from to size(node) always lands in one of the three cases, and each recursive call passes a k from to the size of the subtree it moves into, so the helper never reaches null.

Tracing select(6)

At the root 50, the left subtree has size . Since , the 6th smallest key is in the right subtree, and we have skipped keys: 20, 30, 40, and 50. So we look for the 2nd smallest key under 70. At 70, the left subtree has size . Since , it is 70. The visited nodes are filled:

                 ╭──────╮
                 │█50:9█│
                 ╰───┬──╯
         ┌───────────┴───────────┐
     ╭───┴──╮                ╭───┴──╮
     │ 30:3 │                │█70:5█│
     ╰───┬──╯                ╰───┬──╯
    ┌────┴────┐          ┌───────┴───────┐
╭───┴──╮  ╭───┴──╮   ╭───┴──╮        ╭───┴──╮
│ 20:1 │  │ 40:1 │   │ 60:1 │        │ 80:3 │
╰──────╯  ╰──────╯   ╰──────╯        ╰───┬──╯
                                    ┌────┴────┐
                                ╭───┴──╮  ╭───┴──╮
                                │ 75:1 │  │ 90:1 │
                                ╰──────╯  ╰──────╯

select(6) returns 70, after visiting two nodes.

Cost

Each call does a constant amount of work and moves one level down, so select is , which is on an AVL tree. The iterator version takes steps. Here the cost does not depend on k.