Growing with the Web

AVL tree

Published , updated
Tags:

The AVL tree, named after its inventors Georgy Adelson-Velsky and Evgenii Landis, is a type of self-balancing binary search tree. The tree re-organises itself after every insert and delete so that the tree height is approximately \log n nodes high, allowing search in O(\log n) time. The re-organising does not guarantee a perfectly balanced tree, it is however good enough to guarantee O(\log n) search.

The AVL tree is similar to the self-balancing red-black tree, but instead of augmenting a binary search tree with a colour to help it achieve it’s balancing goals, it keeps track of the height of each node. When the height of the node’s children are not balanced, a rebalancing operation is performed.

This article assumes knowledge of the binary search tree (BST) data structure.

Complexity

Operation Description Complexity
Delete Deletes a node given a key O(\log n)
Insert Inserts a node with an associated key O(\log n)
Search Searches for and returns a node using its key O(\log n)

Properties

An AVL tree is a BST with the following additional property:

  • Each node’s children differ in height by at most 1; |h_l - h_r| \leq 1

Representation

There are two main ways of representing a binary tree. The first is using node objects that have references to their children.

Tree representation

The second is using a regular array and manipulating the index of the node to find its children. The index of the left child of a node is 2i + 1 and the index of the right is 2i + 2 where i is the index of the parent.

Array representation

The index of node's parent can also be retrieved with \lfloor(i - 1) / 2\rfloor.

Height of a node

The height of a node is defined as the maximum number of steps from the node to any of its leaf nodes. This can be defined recursively as h = max(h_l, h_r) + 1, with the base case being a h = 0 for a leaf node.

This is not to be confused with the depth of a node which is the number of steps from the tree’s root.

A helpful implementation detail is to define children that don’t exist as having a height of -1, this allows the balance function defined in the properties section to work on nodes with 0 or 1 children.

Balancing operations

When insertion or deletion occurs, the heights of nodes are updated and a balancing operation occurs. There are four different cases that can occur, each of them can be fixed using either one or two binary tree rotation operations.

Left left case

This case occurs when the height of a node’s left child becomes 2 greater than that of the right child, and the left child is either balanced or left-heavy. It can be fixed using a single right rotate operation on the unbalanced node.

Left right case

This case occurs when the height of a node’s left child becomes 2 greater than that of the right child, and the left child is right-heavy. It can be fixed using a left rotate on the left child, which results in the left left case which is fixed using a right rotate on the unbalanced node.

Right right case

This case occurs when the height of a node’s right child becomes 2 greater than that of the left child, and the right child is either balanced or right-heavy. This is the inverse of the left left case. It can be fixed using a single left rotate operation on the unbalanced node.

Right left case

This case occurs when the height of a node’s right child becomes 2 greater than that of the left child, and the right child is left-heavy. This is the inverse of the left right case. It can be fixed using a right rotate on the right child, which results in the right right case which is fixed using a left rotate on the unbalanced node.

Operations

Delete

Delete starts by performing a regular BST delete. Once the node is deleted, go up the tree updating the heights of each node and rebalancing when required as outlined in the balancing section.

Insert

Insert starts by performing a regular BST insert. Once the node is inserted, go up the tree updating the heights of each node and rebalancing when required as outlined in the balancing section.

Search simply performs a BST search since the AVL maintains BST properties.

Worst case balancing

The worst case balancing of the AVL tree occurs when one child on every node has a height of 1 more than the other child.

Examples of worst case balanced AVL trees of heights 1, 2 and 3

Notice how the example of a tree of height 3 is a node with the left child being the height 2 tree and right child being the height 1 tree. This illustrates the recurrence for the worst case height of a tree.

F_h = F_{h-1} + F_{h-2} + 1

It’s easy to see the worst case height formula in action by highlighting the sub-trees.

A visual example of the recurrence

This recurrence is very similar to the Fibonacci sequence; F_n = F_{n-1} + F_{n-2}. Fortunately, just like the Fibonacci sequence, the number of nodes in a worst case balanced AVL tree increases exponentially. This means that the worst case height is still O(\log n) even in a worst case balanced tree.

Which binary search tree is best?

The AVL tree, red-black tree and splay tree are all self-adjusting binary search trees, so which one is better in which situation? And when is it better to use a regular BST?

A paper by Ben Pfaff of Stanford University performs an in-depth study of the performance characteristics of each tree under various circumstances. Each data structure excels based on runtime patterns in the input and the calling of operations. It comes to the following conclusions:

  • Regular BSTs excel when randomly ordered input can be relied upon.
  • Splay trees excel when data is often inserted in a sorted order and later accesses are sequential or clustered.
  • AVL trees excel when data is often inserted in a sorted order and later accesses are random.
  • Red-black trees excel when data is often inserted in random order but occasional runs of sorted order are expected.

Code

/**
 * Generic implementation of an AVL tree.
 */
public class AVLTree<K extends Comparable<K>> implements BinarySearchTreeInterface<K> {
    /**
     * The root of the tree.
     */
    private AVLTreeNode<K> root;

    /**
     * The size of the tree.
     */
    private int size = 0;

    /**
     * Creates a new {@link AVLTree}.
     */
    public AVLTree() { }

    /** {@inheritDoc} */
    public void insert(K key) {
        root = insert(key, root);
        size++;
    }

    /**
     * Inserts a new node with a specific key into the tree.
     *
     * @param key The key being inserted.
     * @param root The root of the tree to insert in.
     * @return The new tree root.
     */
    private AVLTreeNode<K> insert(K key, AVLTreeNode<K> root) {
        // Perform regular BST insertion
        if (root == null) {
            return new AVLTreeNode<K>(key);
        }

        if (key.compareTo(root.getKey()) < 0) {
            root.setLeft(insert(key, root.getLeft()));
        } else if (key.compareTo(root.getKey()) > 0) {
            root.setRight(insert(key, root.getRight()));
        } else {
            // It's a duplicate so insertion failed, decrement size to make up for it
            size--;
            return root;
        }

        // Update height and rebalance tree
        root.setHeight(Math.max(root.getLeftHeight(), root.getRightHeight()) + 1);
        BalanceState balanceState = getBalanceState(root);

        if (balanceState == BalanceState.UNBALANCED_LEFT) {
            if (key.compareTo(root.getLeft().getKey()) < 0) {
                // Left left case
                root = root.rotateRight();
            } else {
                // Left right case
                root.setLeft(root.getLeft().rotateLeft());
                return root.rotateRight();
            }
        }

        if (balanceState == BalanceState.UNBALANCED_RIGHT) {
            if (key.compareTo(root.getRight().getKey()) > 0) {
                // Right right case
                root = root.rotateLeft();
            } else {
                // Right left case
                root.setRight(root.getRight().rotateRight());
                return root.rotateLeft();
            }
        }

        return root;
    }

    /** {@inheritDoc} */
    public void delete(K key) {
        root = delete(key, root);
        size--;
    }

    /**
     * Deletes a node with a specific key from the tree.
     *
     * @param key The key being deleted.
     * @param root The root of the tree to delete from.
     * @return The new tree root.
     */
    private AVLTreeNode<K> delete(K key, AVLTreeNode<K> root) {
        // Perform regular BST deletion
        if (root == null) {
            size++;
            return root;
        }

        if (key.compareTo(root.getKey()) < 0) {
            // The key to be deleted is in the left sub-tree
            root.setLeft(delete(key, root.getLeft()));
        } else if (key.compareTo(root.getKey()) > 0) {
            // The key to be deleted is in the right sub-tree
            root.setRight(delete(key, root.getRight()));
        } else {
            // root is the node to be deleted
            if (!root.leftExists() && !root.rightExists()) {
                root = null;
            } else if (!root.leftExists() && root.rightExists()) {
                root = root.getRight();
            } else if (root.leftExists() && !root.rightExists()) {
                root = root.getLeft();
            } else {
                // Node has 2 children, get the in-order successor
                AVLTreeNode<K> inOrderSuccessor = minValueNode(root.getRight());
                root.setKey(inOrderSuccessor.getKey());
                root.setRight(delete(inOrderSuccessor.getKey(), root.getRight()));
            }
        }

        if (root == null) {
            return root;
        }

        // Update height and rebalance tree
        root.setHeight(Math.max(root.getLeftHeight(), root.getRightHeight()) + 1);
        BalanceState balanceState = getBalanceState(root);

        if (balanceState == BalanceState.UNBALANCED_LEFT) {
            // Left left case
            if (getBalanceState(root.getLeft()) == BalanceState.BALANCED ||
                    getBalanceState(root.getLeft()) == BalanceState.SLIGHTLY_UNBALANCED_LEFT) {
                return root.rotateRight();
            }
            // Left right case
            if (getBalanceState(root.getLeft()) == BalanceState.SLIGHTLY_UNBALANCED_RIGHT) {
                root.setLeft(root.getLeft().rotateLeft());
                return root.rotateRight();
            }
        }

        // Right right case
        if (balanceState == BalanceState.UNBALANCED_RIGHT) {
            if (getBalanceState(root.getLeft()) == BalanceState.BALANCED ||
                    getBalanceState(root.getLeft()) == BalanceState.SLIGHTLY_UNBALANCED_RIGHT) {
                return root.rotateLeft();
            }
            // Right left case
            if (getBalanceState(root.getLeft()) == BalanceState.SLIGHTLY_UNBALANCED_LEFT) {
                root.setRight(root.getRight().rotateRight());
                return root.rotateLeft();
            }
        }

        return root;
    }

    /** {@inheritDoc} */
    public boolean contains(K key) {
        if (root == null) {
            return false;
        }

        return contains(key, root);
    }

    /**
     * Gets whether a node with a specific key is within the tree.
     *
     * @param key The key being searched for.
     * @param root The root of the tree to search in.
     * @return Whether a node with the key exists.
     */
    private boolean contains(K key, AVLTreeNode<K> root) {
        if (key == root.getKey()) {
            return true;
        }

        if (key.compareTo(root.getKey()) < 0) {
            if (!root.leftExists()) {
                return false;
            }
            return contains(key, root.getLeft());
        }

        if (key.compareTo(root.getKey()) > 0) {
            if (!root.rightExists()) {
                return false;
            }
            return contains(key, root.getRight());
        }

        return false;
    }

    /**
     * @return The minimum key in the tree.
     */
    public K findMinimum() {
        return minValueNode(root).getKey();
    }

    /**
     * Gets the minimum value node, rooted in a particular node.
     *
     * @param root The node to search.
     * @return The node with the minimum key in the tree.
     */
    private AVLTreeNode<K> minValueNode(AVLTreeNode<K> root) {
        AVLTreeNode<K> current = root;
        while (current.leftExists()) {
            current = current.getLeft();
        }
        return current;
    }

    /**
     * @return The maximum key in the tree.
     */
    public K findMaximum() {
        return maxValueNode(root).getKey();
    }

    /**
     * Gets the maximum value node, rooted in a particular node.
     *
     * @param root The node to search.
     * @return The node with the maximum key in the tree.
     */
    private AVLTreeNode<K> maxValueNode(AVLTreeNode<K> root) {
        AVLTreeNode<K> current = root;
        while (current.rightExists()) {
            current = current.getRight();
        }
        return current;
    }

    /** {@inheritDoc} */
    public int size() {
        return size;
    }

    /** {@inheritDoc} */
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * Gets the balance state of a node, indicating whether the left or right sub-trees are
     * unbalanced.
     *
     * @param node The node to get the difference from.
     * @return The {@link BalanceState} of the node.
     */
    private BalanceState getBalanceState(AVLTreeNode<K> node) {
        if (node == null) {
            return BalanceState.BALANCED;
        }
        int heightDifference = node.getLeftHeight() - node.getRightHeight();
        switch (heightDifference) {
            case -2: return BalanceState.UNBALANCED_RIGHT;
            case -1: return BalanceState.SLIGHTLY_UNBALANCED_RIGHT;
            case 1: return BalanceState.SLIGHTLY_UNBALANCED_LEFT;
            case 2: return BalanceState.UNBALANCED_LEFT;
        }
        return BalanceState.BALANCED;
    }

    /**
     * Represents how balanced a node's left and right children are.
     */
    private enum BalanceState {
        UNBALANCED_RIGHT,
        SLIGHTLY_UNBALANCED_RIGHT,
        BALANCED,
        SLIGHTLY_UNBALANCED_LEFT,
        UNBALANCED_LEFT
    }
}



/**
 * Generic implementation of an AVL tree node.
 */
public class AVLTreeNode<K extends Comparable<K>> implements Comparable<AVLTreeNode<K>> {
    /**
     * The height of a left or right child node that doesn't exist.
     */
    private static final int NULL_NODE_HEIGHT = -1;

    /**
     * The left child node.
     */
    private AVLTreeNode<K> left;

    /**
     * The right child node.
     */
    private AVLTreeNode<K> right;

    /**
     * The key of the node.
     */
    private K key;

    /**
     * The height of the node.
     */
    private int height;

    /**
     * Creates a new {@link AVLTreeNode}.
     *
     * @param key The key of the new node.
     */
    public AVLTreeNode(K key) {
        this.key = key;
    }

    /**
     * @return Whether the node has a left child.
     */
    public boolean leftExists() {
        return left != null;
    }

    /**
     * @return Whether the node has a right child.
     */
    public boolean rightExists() {
        return right != null;
    }

    /**
     * @return The key of this node.
     */
    public K getKey() {
        return key;
    }

    /**
     * Sets the key of this node.
     *
     * @param key The new key.
     */
    public void setKey(K key) {
        this.key = key;
    }

    /**
     * @return The left child node.
     */
    public AVLTreeNode<K> getLeft() {
        return left;
    }

    /**
     * Sets the left child node.
     *
     * @param left The new left child node.
     */
    public void setLeft(AVLTreeNode<K> left) {
        this.left = left;
    }

    /**
     * @return The right child node.
     */
    public AVLTreeNode<K> getRight() {
        return right;
    }

    /**
     * Sets the right child node.
     *
     * @param right The new right child node.
     */
    public void setRight(AVLTreeNode<K> right) {
        this.right = right;
    }

    /**
     * @return The height of the node.
     */
    public int getHeight() {
        return height;
    }

    /**
     * Sets the height of the node.
     *
     * @param height The new height of the node.
     */
    public void setHeight(int height) {
        this.height = height;
    }


    /**
     * Performs a right rotate on this node.
     *
     *       b                           a
     *      / \                         / \
     *     a   e -> b.rotateRight() -> c   b
     *    / \                             / \
     *   c   d                           d   e
     *
     * @return The root of the sub-tree; the node where this node used to be.
     */
    public AVLTreeNode<K> rotateRight() {
        AVLTreeNode<K> other = getLeft();
        setLeft(other.getRight());
        other.setRight(this);
        setHeight(Math.max(this.getLeftHeight(), this.getRightHeight()) + 1);
        other.setHeight(Math.max(other.getLeftHeight(), this.getHeight()) + 1);
        return other;
    }

    /**
     * Performs a left rotate on this node.
     *
     *     a                              b
     *    / \                            / \
     *   c   b   -> a.rotateLeft() ->   a   e
     *      / \                        / \
     *     d   e                      c   d
     *
     * @return The root of the sub-tree; the node where this node used to be.
     */
    public AVLTreeNode<K> rotateLeft() {
        AVLTreeNode<K> other = getRight();
        setRight(other.getLeft());
        other.setLeft(this);
        setHeight(Math.max(this.getLeftHeight(), this.getRightHeight()) + 1);
        other.setHeight(Math.max(other.getRightHeight(), this.getHeight()) + 1);
        return other;
    }

    /**
     * Convenience function to get the height of the left child of the node, returning
     * {@link NULL_NODE_HEIGHT} if the node is null.
     *
     * @return The height of the left child, or {@link NULL_NODE_HEIGHT} if it doesn't exist.
     */
    public int getLeftHeight() {
        if (!leftExists()) {
            return NULL_NODE_HEIGHT;
        }
        return getLeft().getHeight();
    }

    /**
     * Convenience function to get the height of the right child of the node, returning
     * {@link NULL_NODE_HEIGHT} if the node is null.
     *
     * @return The height of the right child, or {@link NULL_NODE_HEIGHT} if it doesn't exist.
     */
    public int getRightHeight() {
        if (!rightExists()) {
            return NULL_NODE_HEIGHT;
        }
        return getRight().getHeight();
    }

    /** {@inheritDoc} */
    @Override
    public int compareTo(AVLTreeNode<K> o) {
        return key.compareTo(o.key);
    }
}



/**
 * An interface for data structures that support binary search tree operations.
 */
public interface BinarySearchTreeInterface<K> {
    /**
     * Inserts a new node with a specific key into the tree.
     *
     * @param key The key being inserted.
     */
    void insert(K key);

    /**
     * Deletes a node with a specific key from the tree.
     *
     * @param key The key being deleted.
     */
    void delete(K key);

    /**
     * Gets whether a node with a specific key is within the tree.
     *
     * @param key The key being searched for.
     * @return Whether a node with the key exists.
     */
    boolean contains(K key);

    /**
     * @return The size of the tree.
     */
    int size();

    /**
     * @return Whether the tree is empty; the size of the tree is 0).
     */
    boolean isEmpty();
}
/**
 * Creates a new AVL Tree.
 *
 * @param {function} customCompare An optional custom compare function.
 */
var AvlTree = function (customCompare) {
  this._root = null;
  this._size = 0;

  if (customCompare) {
    this._compare = customCompare;
  }
};

/**
 * Compares two keys with each other.
 *
 * @private
 * @param {Object} a The first key to compare.
 * @param {Object} b The second key to compare.
 * @return {number} -1, 0 or 1 if a < b, a == b or a > b respectively.
 */
AvlTree.prototype._compare = function (a, b) {
  if (a > b) {
    return 1;
  }
  if (a < b) {
    return -1;
  }
  return 0;
};

/**
 * Inserts a new node with a specific key into the tree.
 *
 * @param {Object} key The key being inserted.
 * @param {Object} value The value being inserted.
 */
AvlTree.prototype.insert = function (key, value) {
  this._root = this._insert(key, value, this._root);
  this._size++;
};

/**
 * Inserts a new node with a specific key into the tree.
 *
 * @private
 * @param {Object} key The key being inserted.
 * @param {Object} value The value being inserted.
 * @param {Node} root The root of the tree to insert in.
 * @return {Node} The new tree root.
 */
AvlTree.prototype._insert = function (key, value, root) {
  // Perform regular BST insertion
  if (root === null) {
    return new Node(key, value);
  }

  if (this._compare(key, root.key) < 0) {
    root.left = this._insert(key, value, root.left);
  } else if (this._compare(key, root.key) > 0) {
    root.right = this._insert(key, value, root.right);
  } else {
    // It's a duplicate so insertion failed, decrement size to make up for it
    this._size--;
    return root;
  }

  // Update height and rebalance tree
  root.height = Math.max(root.leftHeight(), root.rightHeight()) + 1;
  var balanceState = getBalanceState(root);

  if (balanceState === BalanceState.UNBALANCED_LEFT) {
    if (this._compare(key, root.left.key) < 0) {
      // Left left case
      root = root.rotateRight();
    } else {
      // Left right case
      root.left = root.left.rotateLeft();
      return root.rotateRight();
    }
  }

  if (balanceState === BalanceState.UNBALANCED_RIGHT) {
    if (this._compare(key, root.right.key) > 0) {
      // Right right case
      root = root.rotateLeft();
    } else {
      // Right left case
      root.right = root.right.rotateRight();
      return root.rotateLeft();
    }
  }

  return root;
};

/**
 * Deletes a node with a specific key from the tree.
 *
 * @param {Object} key The key being deleted.
 */
AvlTree.prototype.delete = function (key) {
  this._root = this._delete(key, this._root);
  this._size--;
};

/**
 * Deletes a node with a specific key from the tree.
 *
 * @private
 * @param {Object} key The key being deleted.
 * @param {Node} root The root of the tree to delete from.
 * @return {Node} The new tree root.
 */
AvlTree.prototype._delete = function (key, root) {
  // Perform regular BST deletion
  if (root === null) {
    this._size++;
    return root;
  }

  if (this._compare(key, root.key) < 0) {
    // The key to be deleted is in the left sub-tree
    root.left = this._delete(key, root.left);
  } else if (this._compare(key, root.key) > 0) {
    // The key to be deleted is in the right sub-tree
    root.right = this._delete(key, root.right);
  } else {
    // root is the node to be deleted
    if (!root.left && !root.right) {
      root = null;
    } else if (!root.left && root.right) {
      root = root.right;
    } else if (root.left && !root.right) {
      root = root.left;
    } else {
      // Node has 2 children, get the in-order successor
      var inOrderSuccessor = minValueNode(root.right);
      root.key = inOrderSuccessor.key;
      root.value = inOrderSuccessor.value;
      root.right = this._delete(inOrderSuccessor.key, root.right);
    }
  }

  if (root === null) {
    return root;
  }

  // Update height and rebalance tree
  root.height = Math.max(root.leftHeight(), root.rightHeight()) + 1;
  var balanceState = getBalanceState(root);

  if (balanceState === BalanceState.UNBALANCED_LEFT) {
    // Left left case
    if (getBalanceState(root.left) === BalanceState.BALANCED ||
        getBalanceState(root.left) === BalanceState.SLIGHTLY_UNBALANCED_LEFT) {
      return root.rotateRight();
    }
    // Left right case
    if (getBalanceState(root.left) === BalanceState.SLIGHTLY_UNBALANCED_RIGHT) {
      root.left = root.left.rotateLeft();
      return root.rotateRight();
    }
  }

  if (balanceState === BalanceState.UNBALANCED_RIGHT) {
    // Right right case
    if (getBalanceState(root.right) === BalanceState.BALANCED ||
        getBalanceState(root.right) === BalanceState.SLIGHTLY_UNBALANCED_RIGHT) {
      return root.rotateLeft();
    }
    // Right left case
    if (getBalanceState(root.right) === BalanceState.SLIGHTLY_UNBALANCED_LEFT) {
      root.right = root.right.rotateRight();
      return root.rotateLeft();
    }
  }

  return root;
};

/**
 * Gets the value of a node within the tree with a specific key.
 *
 * @param {Object} key The key being searched for.
 * @return {Object} The value of the node or null if it doesn't exist.
 */
AvlTree.prototype.get = function (key) {
  if (this._root === null) {
    return null;
  }

  return this._get(key, this._root).value;
};

/**
 * Gets the value of a node within the tree with a specific key.
 *
 * @private
 * @param {Object} key The key being searched for.
 * @param {Node} root The root of the tree to search in.
 * @return {Object} The node or null if it doesn't exist.
 */
AvlTree.prototype._get = function (key, root) {
  var result = this._compare(key, root.key);

  if (result === 0) {
    return root;
  }

  if (result < 0) {
    if (!root.left) {
      return null;
    }
    return this._get(key, root.left);
  }

  if (!root.right) {
    return null;
  }
  return this._get(key, root.right);
};

/**
 * Gets whether a node with a specific key is within the tree.
 *
 * @param {Object} key The key being searched for.
 * @return {boolean} Whether a node with the key exists.
 */
AvlTree.prototype.contains = function (key) {
  if (this._root === null) {
    return false;
  }

  return !!this._get(key, this._root);
};

/**
 * @return {Object} The minimum key in the tree.
 */
AvlTree.prototype.findMinimum = function () {
  return minValueNode(this._root).key;
};

/**
 * Gets the minimum value node, rooted in a particular node.
 *
 * @private
 * @param {Node} root The node to search.
 * @return {Node} The node with the minimum key in the tree.
 */
function minValueNode(root) {
  var current = root;
  while (current.left) {
    current = current.left;
  }
  return current;
}

/**
 * @return {Object} The maximum key in the tree.
 */
AvlTree.prototype.findMaximum = function () {
  return maxValueNode(this._root).key;
};

/**
 * Gets the maximum value node, rooted in a particular node.
 *
 * @private
 * @param {Node} root The node to search.
 * @return {Node} The node with the maximum key in the tree.
 */
function maxValueNode(root) {
  var current = root;
  while (current.right) {
    current = current.right;
  }
  return current;
}

/**
 * @return {number} The size of the tree.
 */
AvlTree.prototype.size = function () {
  return this._size;
};

/**
 * @return {boolean} Whether the tree is empty.
 */
AvlTree.prototype.isEmpty = function () {
  return this._size === 0;
};

/**
 * Represents how balanced a node's left and right children are.
 *
 * @private
 */
var BalanceState = {
  UNBALANCED_RIGHT: 1,
  SLIGHTLY_UNBALANCED_RIGHT: 2,
  BALANCED: 3,
  SLIGHTLY_UNBALANCED_LEFT: 4,
  UNBALANCED_LEFT: 5
};

/**
 * Gets the balance state of a node, indicating whether the left or right
 * sub-trees are unbalanced.
 *
 * @private
 * @param {Node} node The node to get the difference from.
 * @return {BalanceState} The BalanceState of the node.
 */
function getBalanceState(node) {
  var heightDifference = node.leftHeight() - node.rightHeight();
  switch (heightDifference) {
    case -2: return BalanceState.UNBALANCED_RIGHT;
    case -1: return BalanceState.SLIGHTLY_UNBALANCED_RIGHT;
    case 1: return BalanceState.SLIGHTLY_UNBALANCED_LEFT;
    case 2: return BalanceState.UNBALANCED_LEFT;
    default: return BalanceState.BALANCED;
  }
}



/**
 * Creates a new AVL Tree node.
 *
 * @private
 * @param {Object} key The key of the new node.
 * @param {Object} value The value of the new node.
 */
var Node = function (key, value) {
  this.left = null;
  this.right = null;
  this.height = null;
  this.key = key;
  this.value = value;
};

/**
 * Performs a right rotate on this node.
 *
 *       b                           a
 *      / \                         / \
 *     a   e -> b.rotateRight() -> c   b
 *    / \                             / \
 *   c   d                           d   e
 *
 * @return {Node} The root of the sub-tree; the node where this node used to be.
 */
Node.prototype.rotateRight = function () {
  var other = this.left;
  this.left = other.right;
  other.right = this;
  this.height = Math.max(this.leftHeight(), this.rightHeight()) + 1;
  other.height = Math.max(other.leftHeight(), this.height) + 1;
  return other;
};

/**
 * Performs a left rotate on this node.
 *
 *     a                              b
 *    / \                            / \
 *   c   b   -> a.rotateLeft() ->   a   e
 *      / \                        / \
 *     d   e                      c   d
 *
 * @return {Node} The root of the sub-tree; the node where this node used to be.
 */
Node.prototype.rotateLeft = function () {
  var other = this.right;
  this.right = other.left;
  other.left = this;
  this.height = Math.max(this.leftHeight(), this.rightHeight()) + 1;
  other.height = Math.max(other.rightHeight(), this.height) + 1;
  return other;
};

/**
 * Convenience function to get the height of the left child of the node,
 * returning -1 if the node is null.
 *
 * @return {number} The height of the left child, or -1 if it doesn't exist.
 */
Node.prototype.leftHeight = function () {
  if (!this.left) {
    return -1;
  }
  return this.left.height;
};

/**
 * Convenience function to get the height of the right child of the node,
 * returning -1 if the node is null.
 *
 * @return {number} The height of the right child, or -1 if it doesn't exist.
 */
Node.prototype.rightHeight = function () {
  if (!this.right) {
    return -1;
  }
  return this.right.height;
};
// file: avlTree.ts


import { Node } from './node';
import { AvlTree as AvlTreeApi, CompareFunction } from '@tyriar/avl-tree';

/**
 * Represents how balanced a node's left and right children are.
 */
const enum BalanceState {
  /** Right child's height is 2+ greater than left child's height */
  UNBALANCED_RIGHT,
  /** Right child's height is 1 greater than left child's height */
  SLIGHTLY_UNBALANCED_RIGHT,
  /** Left and right children have the same height */
  BALANCED,
  /** Left child's height is 1 greater than right child's height */
  SLIGHTLY_UNBALANCED_LEFT,
  /** Left child's height is 2+ greater than right child's height */
  UNBALANCED_LEFT
}

export class AvlTree<K, V> implements AvlTreeApi<K, V> {
  protected _root: Node<K, V> | null = null;
  private _size: number = 0;
  private _compare: CompareFunction<K>;

  /**
   * Creates a new AVL Tree.
   * @param _compare An optional custom compare function.
   */
  constructor(
    compare?: CompareFunction<K>
  ) {
    this._compare = compare ? compare : this._defaultCompare;
  }

  /**
   * Compares two keys with each other.
   * @param a The first key to compare.
   * @param b The second key to compare.
   * @return -1, 0 or 1 if a < b, a == b or a > b respectively.
   */
  private _defaultCompare(a: K, b: K): number {
    if (a > b) {
      return 1;
    }
    if (a < b) {
      return -1;
    }
    return 0;
  }

  /**
   * Inserts a new node with a specific key into the tree.
   * @param key The key being inserted.
   * @param value The value being inserted.
   */
  public insert(key: K, value?: V): void {
    this._root = this._insert(key, value, this._root);
    this._size++;
  }

  /**
   * Inserts a new node with a specific key into the tree.
   * @param key The key being inserted.
   * @param root The root of the tree to insert in.
   * @return The new tree root.
   */
  private _insert(key: K, value: V | undefined, root: Node<K, V> | null): Node<K, V> {
    // Perform regular BST insertion
    if (root === null) {
      return new Node(key, value);
    }

    if (this._compare(key, root.key) < 0) {
      root.left = this._insert(key, value, root.left);
    } else if (this._compare(key, root.key) > 0) {
      root.right = this._insert(key, value, root.right);
    } else {
      // It's a duplicate so insertion failed, decrement size to make up for it
      this._size--;
      return root;
    }

    // Update height and rebalance tree
    root.height = Math.max(root.leftHeight, root.rightHeight) + 1;
    const balanceState = this._getBalanceState(root);

    if (balanceState === BalanceState.UNBALANCED_LEFT) {
      if (this._compare(key, (<Node<K, V>>root.left).key) < 0) {
        // Left left case
        root = root.rotateRight();
      } else {
        // Left right case
        root.left = (<Node<K, V>>root.left).rotateLeft();
        return root.rotateRight();
      }
    }

    if (balanceState === BalanceState.UNBALANCED_RIGHT) {
      if (this._compare(key, (<Node<K, V>>root.right).key) > 0) {
        // Right right case
        root = root.rotateLeft();
      } else {
        // Right left case
        root.right = (<Node<K, V>>root.right).rotateRight();
        return root.rotateLeft();
      }
    }

    return root;
  }

  /**
   * Deletes a node with a specific key from the tree.
   * @param key The key being deleted.
   */
  public delete(key: K): void {
    this._root = this._delete(key, this._root);
    this._size--;
  }

  /**
   * Deletes a node with a specific key from the tree.
   * @param key The key being deleted.
   * @param root The root of the tree to delete from.
   * @return The new tree root.
   */
  private _delete(key: K, root: Node<K, V> | null): Node<K, V> | null {
    // Perform regular BST deletion
    if (root === null) {
      this._size++;
      return root;
    }

    if (this._compare(key, root.key) < 0) {
      // The key to be deleted is in the left sub-tree
      root.left = this._delete(key, root.left);
    } else if (this._compare(key, root.key) > 0) {
      // The key to be deleted is in the right sub-tree
      root.right = this._delete(key, root.right);
    } else {
      // root is the node to be deleted
      if (!root.left && !root.right) {
        root = null;
      } else if (!root.left && root.right) {
        root = root.right;
      } else if (root.left && !root.right) {
        root = root.left;
      } else {
        // Node has 2 children, get the in-order successor
        const inOrderSuccessor = this._minValueNode(<Node<K, V>>root.right);
        root.key = inOrderSuccessor.key;
        root.value = inOrderSuccessor.value;
        root.right = this._delete(inOrderSuccessor.key, root.right);
      }
    }

    if (root === null) {
      return root;
    }

    // Update height and rebalance tree
    root.height = Math.max(root.leftHeight, root.rightHeight) + 1;
    const balanceState = this._getBalanceState(root);

    if (balanceState === BalanceState.UNBALANCED_LEFT) {
      // Left left case
      if (this._getBalanceState((<Node<K, V>>root.left)) === BalanceState.BALANCED ||
          this._getBalanceState((<Node<K, V>>root.left)) === BalanceState.SLIGHTLY_UNBALANCED_LEFT) {
        return root.rotateRight();
      }
      // Left right case
      // this._getBalanceState(root.left) === BalanceState.SLIGHTLY_UNBALANCED_RIGHT
      root.left = (<Node<K, V>>root.left).rotateLeft();
      return root.rotateRight();
    }

    if (balanceState === BalanceState.UNBALANCED_RIGHT) {
      // Right right case
      if (this._getBalanceState((<Node<K, V>>root.right)) === BalanceState.BALANCED ||
          this._getBalanceState((<Node<K, V>>root.right)) === BalanceState.SLIGHTLY_UNBALANCED_RIGHT) {
        return root.rotateLeft();
      }
      // Right left case
      // this._getBalanceState(root.right) === BalanceState.SLIGHTLY_UNBALANCED_LEFT
      root.right = (<Node<K, V>>root.right).rotateRight();
      return root.rotateLeft();
    }

    return root;
  }

  /**
   * Gets the value of a node within the tree with a specific key.
   * @param key The key being searched for.
   * @return The value of the node (which may be undefined), or null if it
   * doesn't exist.
   */
  public get(key: K): V | undefined | null {
    if (this._root === null) {
      return null;
    }

    const result = this._get(key, this._root);
    if (result === null) {
      return null;
    }

    return result.value;
  }

  /**
   * Gets the value of a node within the tree with a specific key.
   * @param key The key being searched for.
   * @param root The root of the tree to search in.
   * @return The value of the node or null if it doesn't exist.
   */
  private _get(key: K, root: Node<K, V>): Node<K, V> | null {
    const result = this._compare(key, root.key);
    if (result === 0) {
      return root;
    }

    if (result < 0) {
      if (!root.left) {
        return null;
      }
      return this._get(key, root.left);
    }

    if (!root.right) {
      return null;
    }
    return this._get(key, root.right);
  }

  /**
   * Gets whether a node with a specific key is within the tree.
   * @param key The key being searched for.
   * @return Whether a node with the key exists.
   */
  public contains(key: K): boolean {
    if (this._root === null) {
      return false;
    }

    return !!this._get(key, this._root);
  }

  /**
   * @return The minimum key in the tree or null if there are no nodes.
   */
  public findMinimum(): K | null {
    if (this._root === null) {
      return null;
    }
    return this._minValueNode(this._root).key;
  }

  /**
   * Gets the maximum key in the tree or null if there are no nodes.
   */
  public findMaximum(): K | null {
    if (this._root === null) {
      return null;
    }
    return this._maxValueNode(this._root).key;
  }

  /**
   * Gets the size of the tree.
   */
  public get size(): number {
    return this._size;
  }

  /**
   * Gets whether the tree is empty.
   */
  public get isEmpty(): boolean {
    return this._size === 0;
  }

  /**
   * Gets the minimum value node, rooted in a particular node.
   * @param root The node to search.
   * @return The node with the minimum key in the tree.
   */
  private _minValueNode(root: Node<K, V>): Node<K, V> {
    let current = root;
    while (current.left) {
      current = current.left;
    }
    return current;
  }

  /**
   * Gets the maximum value node, rooted in a particular node.
   * @param root The node to search.
   * @return The node with the maximum key in the tree.
   */
  private _maxValueNode(root: Node<K, V>): Node<K, V> {
    let current = root;
    while (current.right) {
      current = current.right;
    }
    return current;
  }

  /**
   * Gets the balance state of a node, indicating whether the left or right
   * sub-trees are unbalanced.
   * @param node The node to get the difference from.
   * @return The BalanceState of the node.
   */
  private _getBalanceState(node: Node<K, V>): BalanceState {
    const heightDifference = node.leftHeight - node.rightHeight;
    switch (heightDifference) {
      case -2: return BalanceState.UNBALANCED_RIGHT;
      case -1: return BalanceState.SLIGHTLY_UNBALANCED_RIGHT;
      case 1: return BalanceState.SLIGHTLY_UNBALANCED_LEFT;
      case 2: return BalanceState.UNBALANCED_LEFT;
      default: return BalanceState.BALANCED;
    }
  }
}



// file: node.ts


export class Node<K, V> {
  public left: Node<K, V> | null = null;
  public right: Node<K, V> | null = null;
  public height: number | null = null;

  /**
   * Creates a new AVL Tree node.
   * @param key The key of the new node.
   * @param value The value of the new node.
   */
  constructor(
    public key: K,
    public value: V | undefined
  ) {
  }

  /**
   * Performs a right rotate on this node.
   * @return The root of the sub-tree; the node where this node used to be.
   * @throws If Node.left is null.
   */
  public rotateRight(): Node<K, V> {
    //     b                           a
    //    / \                         / \
    //   a   e -> b.rotateRight() -> c   b
    //  / \                             / \
    // c   d                           d   e
    const other = <Node<K, V>>this.left;
    this.left = other.right;
    other.right = this;
    this.height = Math.max(this.leftHeight, this.rightHeight) + 1;
    other.height = Math.max(other.leftHeight, this.height) + 1;
    return other;
  }

  /**
   * Performs a left rotate on this node.
   * @return The root of the sub-tree; the node where this node used to be.
   * @throws If Node.right is null.
   */
  public rotateLeft(): Node<K, V> {
    //   a                              b
    //  / \                            / \
    // c   b   -> a.rotateLeft() ->   a   e
    //    / \                        / \
    //   d   e                      c   d
    const other = <Node<K, V>>this.right;
    this.right = other.left;
    other.left = this;
    this.height = Math.max(this.leftHeight, this.rightHeight) + 1;
    other.height = Math.max(other.rightHeight, this.height) + 1;
    return other;
  }

  /**
   * Convenience function to get the height of the left child of the node,
   * returning -1 if the node is null.
   * @return The height of the left child, or -1 if it doesn't exist.
   */
  public get leftHeight(): number {
    if (this.left === null) {
      return -1;
    }
    return this.left.height || 0;
  }

  /**
   * Convenience function to get the height of the right child of the node,
   * returning -1 if the node is null.
   * @return The height of the right child, or -1 if it doesn't exist.
   */
  public get rightHeight(): number {
    if (this.right === null) {
      return -1;
    }
    return this.right.height || 0;
  }
}

References

  • T.H. Cormen, C.E. Leiserson, R.L. Rivest, C. Stein, “Red-Black Trees” in Introduction to Algorithms, 2nd ed., Cambridge, MA: The MIT Press, 2001, ch. 13, sec. 3, p.296
  • B. Pfaff and Stanford University Department of Computer Science (2004, Jun.), Performance Analysis of BSTs in System Software [Online]. Available: http://www.stanford.edu/~blp/papers/libavl.pdf

Textbooks

Here are two CS textbooks I personally recommend; the Algorithm Design Manual (Steven S. Skiena) is a fantastic introduction to data structures and algorithms without getting to deep into the maths side of things, and Introduction to Algorithms (CLRS) which provides a much deeper, math heavy look.

 

Like this article?
Subscribe for more!