The Rebalance Helper

In the AVL tree chapter, we chose a rotation from two balance factors: the unbalanced node’s, and its heavy child’s. When the two are heavy on the same side, a single rotation fixes the node. When they are heavy on opposite sides, it takes a double rotation. A heavy child with a balance factor of $0$ counts as the same side.

The code first needs the balance factor, the height of the left subtree minus the height of the right:

// Assumes node is not null.
private int balanceFactor(Node<K, V> node) {
  return height(node.left) - height(node.right);
}

It uses the height method from the tree chapter, written here as a private helper of TreeMap:

private int height(Node<K, V> node) {
  if (node == null) {
    return -1;
  }
  return 1 + Math.max(height(node.left), height(node.right));
}

Here is the rule in code:

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;
}

rebalance returns the root of the subtree afterward. If it rotates, that is the node the rotation lifts. If not, it is the same node. The checks on the heavy child are strict, < 0 and > 0, so a heavy child with a balance factor of $0$ goes to the single rotation.