Solution: Pruning the Traversal

The iterator version visits every key and keeps the ones in range, so it costs however small the range is. We want the cost to depend on the size of the answer instead.

At a node, everything in the left subtree is smaller than the node’s key, and everything in the right subtree is larger. That lets the traversal skip a whole subtree when it is entirely outside the range.

@Override
public List<K> keysBetween(K lo, K hi) {
  if (lo == null || hi == null || lo.compareTo(hi) > 0) {
    throw new IllegalArgumentException();
  }
  List<K> keys = new ArrayList<>();
  keysBetween(root, lo, hi, keys);
  return keys;
}

// Appends the keys from lo to hi in the subtree rooted at node to keys, in
// ascending order.
private void keysBetween(Node<K, V> node, K lo, K hi, List<K> keys) {
  if (node == null) {
    return;
  }
  if (node.key.compareTo(lo) > 0) {
    keysBetween(node.left, lo, hi, keys);
  }
  if (node.key.compareTo(lo) >= 0 && node.key.compareTo(hi) <= 0) {
    keys.add(node.key);
  }
  if (node.key.compareTo(hi) < 0) {
    keysBetween(node.right, lo, hi, keys);
  }
}

The three checks are independent:

  • Go left only when the node’s key is strictly greater than lo. Otherwise the whole left subtree is below the range.
  • Add the node’s key when it is from lo to hi.
  • Go right only when the node’s key is strictly less than hi. Otherwise the whole right subtree is above the range.

The order is left, then the node, then right. That is an in-order traversal with two kinds of branch pruned, so the keys still come out in ascending order.

Tracing the query

Run keysBetween(55, 78) on the example tree from the problem.

At the root 50: it is not in range, since . Because 50 is not greater than lo, every key in its left subtree is below 55, so we skip the whole left subtree: 30, 20, and 40 are never visited. Since , we go right.

At 70, we go left first. Its left child 60 is in range, so we add 60. It is a leaf. Back at 70: it is in range too, so we add 70, then go right.

At 80: it is not in range, since . Since , we go left and reach 75, which is in range, so we add 75. Back at 80: because 80 is not less than hi, we skip its right subtree, so 90 is never visited.

The result, in order, is [60, 70, 75].

Here is the finished traversal. The filled nodes are the ones added, and the light nodes are the ones never visited:

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

Cost

The recursion visits the nodes in the range, plus the nodes along two paths: one down toward lo and one down toward hi. Each path is at most long, so the query is . On an AVL tree the height is , which makes it .