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
kis at mostleftSize, the k-th smallest key is in the left subtree, and it is the k-th smallest there too. - If
kisleftSize + 1, it is this node’s key. - Otherwise it is in the right subtree. When we go right, we skip the
leftSizekeys 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.