16. AVL Trees

Written by Kelvin Lau

In the previous chapter, you learned about the O(log n) performance characteristics of the binary search tree. However, you also learned that unbalanced trees can deteriorate the performance of the tree, all the way down to O(n). In 1962, Georgy Adelson-Velsky and Evgenii Landis came up with the first self-balancing binary search tree: The AVL Tree. In this chapter, you’ll dig deeper into how the balance of a binary search tree can impact performance and implement the AVL tree from scratch!

Understanding balance

A balanced tree is the key to optimizing the performance of the binary search tree. In this section, you’ll learn about the three main states of balance.

Perfect balance

The ideal form of a binary search tree is the perfectly balanced state. In technical terms, this means every level of the tree is filled with nodes, from top to bottom.

Not only is the tree perfectly symmetrical, the nodes at the bottom level are completely filled. This is the requirement for being perfectly balanced.

“Good-enough” balance

Although achieving perfect balance is ideal, it is rarely possible. A perfectly balanced tree must contain the exact number of nodes to fill every level to the bottom, so it can only be perfect with a particular number of elements.

For example, a tree with 1, 3 or 7 nodes can be perfectly balanced, but a tree with 2, 4, 5 or 6 cannot be perfectly balanced since the last level of the tree will not be filled.

The definition of a balanced tree is that every level of the tree must be filled, except for the bottom level. In most cases of binary trees, this is the best you can do.

Unbalanced

Finally, there’s the unbalanced state. Binary search trees in this state suffer from various levels of performance loss, depending on the degree of imbalance.

Keeping the tree balanced gives the find, insert and remove operations an O(log n) time complexity. AVL trees maintain balance by adjusting the structure of the tree when the tree becomes unbalanced. You’ll learn how this works as you progress through the chapter.

Implementation

Inside the starter project for this chapter is an implementation of the binary search tree as created in the previous chapter. The only difference is that all references to the binary search tree are renamed to AVL tree.

Binary search trees and AVL trees share much of the same implementation; in fact, all that you’ll add is the balancing component. Open the starter project to begin.

Measuring balance

To keep a binary tree balanced, you’ll need a way to measure the balance of the tree. The AVL tree achieves this with a height property in each node. In tree-speak, the height of a node is the longest distance from the current node to a leaf node:

Open the starter playground for this chapter and add the following property to AVLNode in the compiled sources folder:

public var height = 0

You’ll use the relative heights of a node’s children to determine whether a particular node is balanced. The height of the left and right children of each node must differ at most by 1. This number is known as the balance factor.

Write the following just below the height property of AVLNode:

public var balanceFactor: Int {
  leftHeight - rightHeight
}

public var leftHeight: Int {
  leftChild?.height ?? -1
}

public var rightHeight: Int {
  rightChild?.height ?? -1
}

The balanceFactor computes the height difference of the left and right child. If a particular child is nil, its height is considered to be -1.

Here’s an example of an AVL tree:

The diagram shows a balanced tree — all levels except the bottom one are filled. The numbers to the right of the node represent the height of each node, while the numbers to the left represent the balanceFactor.

Here’s an updated diagram with 40 inserted:

Inserting 40 into the tree turns it into an unbalanced tree. Notice how the balanceFactor changes. A balanceFactor of 2 or -2 or something more extreme indicates of an unbalanced tree. By checking after each insertion or deletion, though, you can guarantee that it is never more extreme than a magnitude of two.

Although more than one node may have a bad balancing factor, you only need to perform the balancing procedure on the bottom-most node containing the invalid balance factor: the node containing 25.

That’s where rotations come in.

Rotations

The procedures used to balance a binary search tree are known as rotations. There are four rotations in total for the four different ways that a tree can become unbalanced. These are known as left rotation, left-right rotation, right rotation and right-left rotation.

Left rotation

The imbalance caused by inserting 40 into the tree can be solved by a left rotation. A generic left rotation of node x looks like this:

Before going into specifics, there are two takeaways from this before-and-after comparison:

• In-order traversal for these nodes remains the same.
• The depth of the tree is reduced by one level after the rotation.

Add the following method to AVLTree, just below insert(from:value:):

private func leftRotate(_ node: AVLNode<Element>)
  -> AVLNode<Element> {

  // 1
  let pivot = node.rightChild!
  // 2
  node.rightChild = pivot.leftChild
  // 3
  pivot.leftChild = node
  // 4
  node.height = max(node.leftHeight, node.rightHeight) + 1
  pivot.height = max(pivot.leftHeight, pivot.rightHeight) + 1
  // 5
  return pivot
}

Here are the steps needed to perform a left rotation:

1. The right child is chosen as the pivot. This node will replace the rotated node as the root of the subtree (it will move up a level).

2. The node to be rotated will become the left child of the pivot (it moves down a level). This means that the current left child of the pivot must be moved elsewhere.

   In the generic example shown in the earlier image, this is node b. Because b is smaller than y but greater than x, it can replace y as the right child of x. So you update the rotated node’s rightChild to the pivot’s leftChild.

3. The pivot’s leftChild can now be set to the rotated node.

4. You update the heights of the rotated node and the pivot.

5. Finally, you return the pivot so that it can replace the rotated node in the tree.

Here are the before-and-after effects of the left rotation of 25 from the previous example:

Right rotation

Right rotation is the symmetrical opposite of left rotation. When a series of left children is causing an imbalance, it’s time for a right rotation.

A generic right rotation of node x looks like this:

To implement this, add the following code just after leftRotate:

private func rightRotate(_ node: AVLNode<Element>)
  -> AVLNode<Element> {

  let pivot = node.leftChild!
  node.leftChild = pivot.rightChild
  pivot.rightChild = node
  node.height = max(node.leftHeight, node.rightHeight) + 1
  pivot.height = max(pivot.leftHeight, pivot.rightHeight) + 1
  return pivot
}

This algorithm is nearly identical to the implementation of leftRotate, except the references to the left and right children are swapped.

Right-left rotation

You may have noticed that the left and right rotations balance nodes that are all left children or all right children. Consider the case in which 36 is inserted into the original example tree.

The right-left rotation:

Doing a left rotation, in this case, won’t result in a balanced tree. The way to handle cases like this is to perform a right rotation on the right child before doing the left rotation. Here’s what the procedure looks like:

1. You apply a right rotation to 37.
2. Now that nodes 25, 36 and 37 are all right children; you can apply a left rotation to balance the tree.

Add the following code just after rightRotate:

private func rightLeftRotate(_ node: AVLNode<Element>)
  -> AVLNode<Element> {

  guard let rightChild = node.rightChild else {
    return node
  }
  node.rightChild = rightRotate(rightChild)
  return leftRotate(node)
}

Don’t worry just yet about when to call this. You’ll get to that in a second. You first need to handle the last case, left-right rotation.

Left-right rotation

Left-right rotation is the symmetrical opposite of the right-left rotation. Here’s an example:

1. You apply a left rotation to node 10.
2. Now that nodes 25, 15 and 10 are all left children; you can apply a right rotation to balance the tree.

Add the following code just after rightLeftRotate:

private func leftRightRotate(_ node: AVLNode<Element>)
  -> AVLNode<Element> {

  guard let leftChild = node.leftChild else {
    return node
  }
  node.leftChild = leftRotate(leftChild)
  return rightRotate(node)
}

That’s it for rotations. Next, you’ll figure out when to apply these rotations at the correct location.

Balance

The next task is to design a method that uses balanceFactor to decide whether a node requires balancing or not. Write the following method below leftRightRotate:

private func balanced(_ node: AVLNode<Element>)
  -> AVLNode<Element> {

  switch node.balanceFactor {
  case 2:
    // ...
  case -2:
    // ...
  default:
    return node
  }
}

There are three cases to consider.

1. A balanceFactor of 2 suggests that the left child is “heavier” (contains more nodes) than the right child. This means that you want to use either right or left-right rotations.
2. A balanceFactor of -2 suggests that the right child is heavier than the left child. This means that you want to use either left or right-left rotations.
3. The default case suggests that the particular node is balanced. There’s nothing to do here except to return the node.

The sign of the balanceFactor can be used to determine if a single or double rotation is required:

Update the balanced function to the following:

private func balanced(_ node: AVLNode<Element>)
  -> AVLNode<Element> {

  switch node.balanceFactor {
  case 2:
    if let leftChild = node.leftChild,
           leftChild.balanceFactor == -1 {
      return leftRightRotate(node)
    } else {
      return rightRotate(node)
    }
  case -2:
    if let rightChild = node.rightChild,
           rightChild.balanceFactor == 1 {
      return rightLeftRotate(node)
    } else {
      return leftRotate(node)
    }
  default:
    return node
  }
}

balanced inspects the balanceFactor to determine the proper course of action. All that’s left is to call balance at the proper spot.

Revisiting insertion

You’ve already done the majority of the work. The remainder is fairly straightforward. Update insert(from:value:) to the following:

private func insert(from node: AVLNode<Element>?,
                    value: Element) -> AVLNode<Element> {
  guard let node = node else {
    return AVLNode(value: value)
  }
  if value < node.value {
    node.leftChild = insert(from: node.leftChild, value: value)
  } else {
    node.rightChild = insert(from: node.rightChild, value: value)
  }
  let balancedNode = balanced(node)
  balancedNode.height = max(balancedNode.leftHeight, balancedNode.rightHeight) + 1
  return balancedNode
}

Instead of returning the node directly after inserting, you pass it into balanced. Passing it ensures every node in the call stack is checked for balancing issues. You also update the node’s height.

That’s all there is to it! Head into the playground page and test it out. Add the following to the playground:

example(of: "repeated insertions in sequence") {
  var tree = AVLTree<Int>()
  for i in 0..<15 {
    tree.insert(i)
  }
  print(tree)
}

You should see the following output in the console:

---Example of: repeated insertions in sequence---
  ┌──14
 ┌──13
 │ └──12
┌──11
│ │ ┌──10
│ └──9
│  └──8
7
│  ┌──6
│ ┌──5
│ │ └──4
└──3
 │ ┌──2
 └──1
  └──0

Take a moment to appreciate the uniform spread of the nodes. If the rotations weren’t applied, this would have become a long, unbalanced link of right children.

Revisiting remove

Retrofitting the remove operation for self-balancing is just as easy as fixing insert. In AVLTree, find remove and replace the final return statement with the following:

let balancedNode = balanced(node)
balancedNode.height = max(balancedNode.leftHeight, balancedNode.rightHeight) + 1
return balancedNode

Head back to the playground page and add the following code at the bottom of the file:

example(of: "removing a value") {
  var tree = AVLTree<Int>()
  tree.insert(15)
  tree.insert(10)
  tree.insert(16)
  tree.insert(18)
  print(tree)
  tree.remove(10)
  print(tree)
}

You should see the following console output:

---Example of: removing a value---
 ┌──18
┌──16
│ └──nil
15
└──10

┌──18
16
└──15

Removing 10 caused a left rotation on 15. Feel free to try out a few more test cases of your own.

Whew! The AVL tree is the culmination of your search for the ultimate binary search tree. The self-balancing property guarantees that the insert and remove operations function at optimal performance with an O(log n) time complexity.

Key points

• A self-balancing tree avoids performance degradation by performing a balancing procedure whenever you add or remove elements in the tree.
• AVL trees preserve balance by readjusting parts of the tree when the tree is no longer balanced.
• Balance is achieved by four types of tree rotations on node insertion and removal.

Where to go from here?

While AVL trees were the first self-balancing implementations of a BST, others, such as the red-black tree and splay tree, have since joined the party. If you’re interested, you check those out in the raywenderlich.com Swift Algorithm Club. Find them at at: https://github.com/raywenderlich/swift-algorithm-club/tree/master/Red-Black%20Tree and https://github.com/raywenderlich/swift-algorithm-club/tree/master/Splay%20Tree respectively.