Rotations and Rebalancing in Code

The questions in this section use this outline of TreeMap. The three helpers are already written. Every helper you write returns the new root of the subtree it was given.

public class TreeMap<K extends Comparable<K>, V> implements Map<K, V> {
  private Node<K, V> root;
  private int size;

  private static class Node<K, V> {
    K key;
    V value;
    Node<K, V> left;
    Node<K, V> right;
    int height;  // cached height of this subtree

    // Sets key and value, with no children and height 0.
    Node(K key, V value) { ... }
  }

  // Returns node.height, or -1 if node is null.
  private int height(Node<K, V> node) { ... }

  // Sets node.height from its children's cached heights.
  private void updateHeight(Node<K, V> node) { ... }

  // Returns height(node.left) - height(node.right).
  private int balanceFactor(Node<K, V> node) { ... }
}

private Node<K, V> rightRotation(Node<K, V> root) {
  Node<K, V> child = root.left;
  root.left = child.right;
  child.right = root;
  updateHeight(root);
  updateHeight(child);
  return child;
}

private Node<K, V> leftRightRotation(Node<K, V> root) {
  root.left = leftRotation(root.left);
  return rightRotation(root);
}

private Node<K, V> rebalance(Node<K, V> node) {
  int bf = balanceFactor(node);
  if (bf > 1) {
    if (balanceFactor(node.left) < 0) {
      return leftRightRotation(node);
    }
    return rightRotation(node);
  }
  if (bf < -1) {
    if (balanceFactor(node.right) > 0) {
      return rightLeftRotation(node);
    }
    return leftRotation(node);
  }
  return node;
}