Sunday, 9 December 2012

Data structure: Red-black tree

The red-black tree is a type of self-balancing binary search tree that assigns a colour of red or black to each node. On every insert or delete, the tree re-organises itself so that it 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.

Insert and delete are also performed in \(O(\log n)\) time. The 'fixup' operations where the balancing occurs after insert and delete have quite complex implementations as you will see below. This is because we need the properties of the red-black tree to hold otherwise it may not be balanced.

Red-black tree example

Complexity

Operation Average case Worst case
Delete \(O(\log n)\) \(O(\log n)\)
Insert \(O(\log n)\) \(O(\log n)\)
Search \(O(\log n)\) \(O(\log n)\)
Space \(O(n)\) \(O(n)\)

Representation

See binary search tree for information on representing a binary tree.

Properties

Here are the set of rules that a red-black tree must follow:

  • Each node is either red or black
  • The root node is black
  • All leaf nodes (nil) are black
  • Both children of every red node are black
  • For each node, all paths from the node to descendant leaves contain the same number of black nodes

Operations

Rotation

Performing a left rotate and a right rotate on nodes is an important operation used in both delete and insert operations. Here is an illustration of the process:

Rotation

Delete(a)

Delete performs a slightly modified binary search tree delete and then performs a fix-up function. Here are the cases that need to be fixed up if they occur.

  1. The deleted node a's sibling is red
  2. The deleted node a's sibling b is black and both of b's children and black
  3. The deleted node a's sibling b is black, b's left child is red and a's right child is black
  4. The deleted node a's sibling b is black and b's right child is red

Insert(a)

Insert performs a slightly modified binary search tree insert and then performs a fix-up function. Here are the cases that need to be fixed up if they occur.

  1. The inserted node a's uncle is red
  2. The inserted node a's uncle is black and a is a right child
  3. The inserted node a's uncle is black and a is a left child

Search(n)

For search we use the search operation on the regular binary search tree. Only now instead of having a worst case of \(O(n)\), it's \(O(\log n)\) as we balance the tree.

Code

View on GitHub

public class RedBlackTree {
    private TreeNode root;

    public RedBlackTree() { }

    public void insert(int key) {
        TreeNode parent = null;
        TreeNode node = root;
        while (node != null && !node.isNilNode()) {
            parent = node;
            if (key < parent.getKey())
                node = parent.getLeft();
            else
                node = parent.getRight();
        }
        if (parent == null) {
            node = new TreeNode(key, null);
            root = node;
        } else {
            node.setParent(parent);
            node.setKey(key);
            node.setNilNode(false);
            node.setColor(TreeNode.Color.RED);
        }
        node.setColor(TreeNode.Color.RED);
        insertFixup(node);
    }

    private void insertFixup(TreeNode node) {
        while (node.getParent() != null &&
               node.getGrandparent() != null &&
               node.getParent().getColor() == TreeNode.Color.RED) {

            if (node.getParent() == node.getGrandparent().getLeft()) {
                TreeNode uncle = node.getGrandparent().getRight();
                if (uncle.getColor() == TreeNode.Color.RED) {
                    node.getParent().setColor(TreeNode.Color.BLACK);
                    uncle.setColor(TreeNode.Color.BLACK);
                    node = node.getGrandparent();
                    node.setColor(TreeNode.Color.RED);
                } else {
                    if (node == node.getParent().getRight()) {
                        node = node.getParent();
                        rotateLeft(node);
                    }
                    node.getParent().setColor(TreeNode.Color.BLACK);
                    node.getGrandparent().setColor(TreeNode.Color.RED);
                    rotateRight(node.getGrandparent());
                }
            } else if (node.getParent() == node.getGrandparent().getRight()) {
                TreeNode uncle = node.getGrandparent().getLeft();
                if (uncle.getColor() == TreeNode.Color.RED) {
                    node.getParent().setColor(TreeNode.Color.BLACK);
                    uncle.setColor(TreeNode.Color.BLACK);
                    node = node.getGrandparent();
                    node.setColor(TreeNode.Color.RED);
                } else {
                    if (node == node.getParent().getLeft()) {
                        node = node.getParent();
                        rotateRight(node);
                    }
                    node.getParent().setColor(TreeNode.Color.BLACK);
                    node.getGrandparent().setColor(TreeNode.Color.RED);
                    rotateLeft(node.getGrandparent());
                }
            }
        }
        root.setColor(TreeNode.Color.BLACK);
    }

    private void rotateLeft(TreeNode x) {
        TreeNode y = x.getRight();
        x.setRight(y.getLeft());
        if (y.getLeft() != null)
            y.getLeft().setParent(x);
        y.setParent(x.getParent());
        if (x.getParent() == null)
            root = y;
        else {
            if (x == x.getParent().getLeft())
                x.getParent().setLeft(y);
            else
                x.getParent().setRight(y);
        }
        y.setLeft(x);
        x.setParent(y);
    }

    private void rotateRight(TreeNode x) {
        TreeNode y = x.getLeft();
        x.setLeft(y.getRight());
        if (y.getRight() != null)
            y.getRight().setParent(x);
        y.setParent(x.getParent());
        if (x.getParent() == null)
            root = y;
        else {
            if (x == x.getParent().getLeft())
                x.getParent().setLeft(y);
            else
                x.getParent().setRight(y);
        }
        y.setRight(x);
        x.setParent(y);
    }

    public void delete(int key) {
        TreeNode node = search(key);
        TreeNode y, x;
        if (node.getLeft().isNilNode() || node.getRight().isNilNode())
            y = node;
        else
            y = treeSuccessor(node);
        if (y.getLeft() != null && !y.getLeft().isNilNode())
            x = y.getLeft();
        else
            x = y.getRight();
        x.setParent(y.getParent());
        if (y.getParent() == null)
            root = x;
        else {
            if (y == y.getParent().getLeft())
                y.getParent().setLeft(x);
            else
                y.getParent().setRight(x);
        }
        if (y != node)
            node.setKey(y.getKey());
        if (y.getColor() == TreeNode.Color.BLACK)
            deleteFixup(x);
    }

    private void deleteFixup(TreeNode node) {
        while (node != root && node.getColor() == TreeNode.Color.BLACK) {
            if (node == node.getParent().getLeft()) {
                TreeNode w = node.getParent().getRight();
                if (w.getColor() == TreeNode.Color.RED) {
                    w.setColor(TreeNode.Color.BLACK);
                    node.getParent().setColor(TreeNode.Color.RED);
                    rotateLeft(node.getParent());
                }
                if (w.getLeft().getColor() == TreeNode.Color.BLACK &&
                    w.getRight().getColor() == TreeNode.Color.BLACK) {

                    w.setColor(TreeNode.Color.RED);
                    node = node.getParent();
                } else  {
                    if (w.getRight().getColor() == TreeNode.Color.BLACK) {
                        w.getLeft().setColor(TreeNode.Color.BLACK);
                        w.setColor(TreeNode.Color.RED);
                        rotateRight(w);
                        w = node.getParent().getRight();
                    }
                    w.setColor(node.getParent().getColor());
                    node.getParent().setColor(TreeNode.Color.BLACK);
                    w.getRight().setColor(TreeNode.Color.BLACK);
                    rotateLeft(node.getParent());
                    node = root;
                }
            } else {
                TreeNode w = node.getParent().getLeft();
                if (w.getColor() == TreeNode.Color.RED) {
                    w.setColor(TreeNode.Color.BLACK);
                    node.getParent().setColor(TreeNode.Color.RED);
                    rotateRight(node.getParent());
                }
                if (w.getRight().getColor() == TreeNode.Color.BLACK &&
                    w.getLeft().getColor() == TreeNode.Color.BLACK) {

                    w.setColor(TreeNode.Color.RED);
                    node = node.getParent();
                } else  {
                    if (w.getLeft().getColor() == TreeNode.Color.BLACK) {
                        w.getRight().setColor(TreeNode.Color.BLACK);
                        w.setColor(TreeNode.Color.RED);
                        rotateLeft(w);
                        w = node.getParent().getLeft();
                    }
                    w.setColor(node.getParent().getColor());
                    node.getParent().setColor(TreeNode.Color.BLACK);
                    w.getLeft().setColor(TreeNode.Color.BLACK);
                    rotateRight(node.getParent());
                    node = root;
                }
            }
        }
        node.setColor(TreeNode.Color.BLACK);
    }

    private TreeNode treeSuccessor(TreeNode node) {
        if (node.getRight() != null && !node.isNilNode())
            return treeMinimum(node.getRight());
        TreeNode successor = node.getParent();
        while (successor != null && !successor.isNilNode() && node == successor) {
            node = successor;
            successor = node.getParent();
        }
        return successor;
    }

    private TreeNode treeMinimum(TreeNode node) {
        while (!node.getLeft().isNilNode() && !node.isNilNode())
            node = node.getLeft();
        return node;
    }

    public TreeNode search(int key) {
        if (root == null)
            return null;

        return search(key, root);
      }

    public TreeNode search(int key, TreeNode node) {
        if (key == node.getKey())
            return node;

        if (key < node.getKey()) {
            if (!node.leftExists())
                return null;
            return search(key, node.getLeft());
        }

        if (key > node.getKey()) {
            if (!node.rightExists())
                return null;
            return search(key, node.getRight());
        }

        return null;
    }

    public String toString() {
        if (root == null)
            return "(empty)";
        return root.toString();
    }
}

Here is the TreeNode class. There is a little extra logic in it on top of the binary search tree implementation to fulfil the property "All leaf nodes (nil) are black".

public class TreeNode {
    private final String nullNodeString = "_B";
    private TreeNode left;
    private TreeNode right;
    private TreeNode parent;

    private int key;
    private boolean isNilNode;
    private Color color;

    public TreeNode(int key, TreeNode parent) {
        this.key = key;
        this.parent = parent;
        this.color = Color.RED;
        this.setNilNode(false);
    }

    // Constructor for nil leaf node
    private TreeNode(TreeNode parent) {
        this.parent = parent;
        this.color = Color.BLACK;
        this.setNilNode(true);
    }

    public String toString() {
        if (isNilNode)
            return nullNodeString;
        return key + getColorCode() + " : { " +
                (leftExists() ? left.toString() : nullNodeString) + " , " +
                (rightExists() ? right.toString() : nullNodeString) + " }";
    }

    private String getColorCode() {
        if (color == Color.BLACK)
            return "B";
        return "R";
    }

    public boolean leftExists() {
        return left != null;
    }

    public boolean rightExists() {
        return right != null;
    }

    public int getKey() {
        return key;
    }

    public void setKey(int key) {
        this.key = key;
    }

    public TreeNode getLeft() {
        // Create nil leaf nodes lazily
        if (left == null)
            left = new TreeNode(this);
        return left;
    }

    public void setLeft(TreeNode left) {
        this.left = left;
    }

    public TreeNode getRight() {
        // Create nil leaf nodes lazily
        if (right == null)
            right = new TreeNode(this);
        return right;
    }

    public void setRight(TreeNode right) {
        this.right = right;
    }

    public TreeNode getParent() {
        return parent;
    }

    public TreeNode getGrandparent() {
        if (parent != null && parent.getParent() != null)
            return parent.getParent();
        return null;
    }

    public void setParent(TreeNode parent) {
        this.parent = parent;
    }

    public Color getColor() {
        return color;
    }

    public void setColor(Color color) {
        this.color = color;
    }

    public boolean isNilNode() {
        return isNilNode;
    }

    public void setNilNode(boolean isNilNode) {
        this.isNilNode = isNilNode;
    }

    public enum Color {
        BLACK,
        RED
    }
}