put, First Version

put does two jobs. If the key is new, it adds a node and returns true. If the key is already there, it replaces that key’s value and returns false.

Here is the recursive helper for put:

private Node<K, V> put(Node<K, V> node, K key, V value) {
  if (node == null) {
    return new Node<>(key, value);
  }

  int cmp = key.compareTo(node.key);
  if (cmp < 0) {
    node.left = put(node.left, key, value);
  } else if (cmp > 0) {
    node.right = put(node.right, key, value);
  } else {
    node.value = value;  // no new node, so the shape did not change
    return node;
  }
  return rebalance(node);
}

When it reaches a matching key, it replaces the value and returns the node as it is. Only the new-node case changes the tree’s shape, so only that case goes on to rebalance.

The public method has to return whether the key was new and keep size correct. For that it needs to know whether the key was there before the recursion runs. So it checks first, then calls the helper either way:

@Override
public boolean put(K key, V value) {
  if (key == null || value == null) {
    throw new IllegalArgumentException();
  }
  boolean isNew = !containsKey(key);
  root = put(root, key, value);
  if (isNew) {
    size++;
  }
  return isNew;
}

The assignment root = put(root, key, value) stores the new root of the whole tree, in case a rotation changed it. This costs two searches down the tree, one in containsKey and one in the helper. We could have the recursion report whether it added a node and skip the first search. I keep the two-step version because it is clearer, and asymptotically it makes no difference.

To see it work, put the key 4 into the tree below. This is the insertion the AVL tree chapter used to show that one repair is enough: 11 is the first node that becomes unbalanced, and a right rotation at 11 fixes it. The illustrations show only the keys.

         ╭────╮
         │ 14 │
         ╰──┬─╯
      ┌─────┴─────┐
   ╭──┴─╮      ╭──┴─╮
   │ 11 │      │ 17 │
   ╰──┬─╯      ╰──┬─╯
  ┌───┘           └───┐
╭─┴─╮              ╭──┴─╮
│ 7 │              │ 19 │
╰───╯              ╰────╯

The base case

At 14, 11, and 7, the key 4 is smaller, so each call goes left and waits for the call below it. The call on 7.left hits the base case: it creates the new node and returns it.

call stack

        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
top ──▶ ┃ put(node=null)                         ┃
        ┃   return new Node(4)  (the empty slot) ┃
        ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
        ┌────────────────────────────────────────┐
        │ put(node=7)                            │
        │   waiting for put(7.left)              │
        ├────────────────────────────────────────┤
        │ put(node=11)                           │
        │   waiting for put(11.left)             │
        ├────────────────────────────────────────┤
        │ put(node=14)                           │
        │   waiting for put(14.left)             │
        └────────────────────────────────────────┘

tree

         ╭────╮
         │ 14 │
         ╰──┬─╯
      ┌─────┴─────┐
   ╭──┴─╮      ╭──┴─╮
   │ 11 │      │ 17 │
   ╰──┬─╯      ╰──┬─╯
  ┌───┘           └───┐
╭─┴─╮              ╭──┴─╮
│█7█│              │ 19 │
╰───╯              ╰────╯

Unwinding, with a rotation

put(7) receives the new node. It runs 7.left = <new node>, then rebalance(7). The balance factor of 7 is $1$, which is within $-1$ to $1$, so rebalance returns 7 unchanged.

call stack

        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
top ──▶ ┃ put(node=11)               ┃ ◀── receives 7
        ┃   waiting for put(11.left) ┃
        ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛
        ┌────────────────────────────┐
        │ put(node=14)               │
        │   waiting for put(14.left) │
        └────────────────────────────┘

tree

              ╭────╮
              │ 14 │
              ╰──┬─╯
          ┌──────┴──────┐
       ╭──┴─╮        ╭──┴─╮
       │ 11 │        │ 17 │
       ╰──┬─╯        ╰──┬─╯
      ┌───┘             └───┐
    ╭─┴─╮                ╭──┴─╮
    │█7█│                │ 19 │
    ╰─┬─╯                ╰────╯
  ┌───┘
╭─┴─╮
│ 4 │
╰───╯

put(11) receives 7. It runs 11.left = 7, then rebalance(11). This time the balance factor of 11 is $2$, the left-left shape, so rebalance does a single right rotation and returns 7 instead of 11.

call stack

        ┏━━━━━━━━━━━━━━━━━━━━━━━━━━━━┓
top ──▶ ┃ put(node=14)               ┃ ◀── receives 7
        ┃   waiting for put(14.left) ┃
        ┗━━━━━━━━━━━━━━━━━━━━━━━━━━━━┛

returned by rebalance(11)

     ╭───╮
     │░7░│
     ╰─┬─╯
  ┌────┴────┐
╭─┴─╮    ╭──┴─╮
│ 4 │    │░11░│
╰───╯    ╰────╯

put(14) receives 7 and stores it with 14.left = 7. Because of the rotation at 11, the left child of 14 is now 7, not 11. Then rebalance(14) finds a balance factor of $0$ and returns 14. The public method stores it into root.

call stack

        (empty — all calls have returned)

tree

             ╭────╮
             │█14█│
             ╰──┬─╯
       ┌────────┴────────┐
     ╭─┴─╮            ╭──┴─╮
     │ 7 │            │ 17 │
     ╰─┬─╯            ╰──┬─╯
  ┌────┴────┐            └───┐
╭─┴─╮    ╭──┴─╮           ╭──┴─╮
│ 4 │    │ 11 │           │ 19 │
╰───╯    ╰────╯           ╰────╯