How to Implement a Red Black Tree

import java.util.ArrayList;


public class RedBlackTree<K extends Comparable<K>, V> {

    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private Node root;

    private int size;

    public static void main(String[] args) {
        ArrayList<Integer> words = new ArrayList<>();

        for (int i = 0; i < 10000; i++) {
            words.add(i);
        }

        RedBlackTree<Integer, Integer> tree = new RedBlackTree<>();

        for (Integer word : words) {
            if (!tree.contains(word)) {
                tree.insert(word, 0);
            }

            tree.set(word, tree.search(word).value + 1);
        }

        System.out.println("Height: " + tree.getHeight(tree.root));

        System.out.println(tree.getMin().key + ": " + tree.getMin().value);

        tree.remove(tree.getMin().key);

        System.out.println(tree.getMin().key + ": " + tree.getMin().value);

        System.out.println("Is valid BST? " + tree.isBST());
    }

    /**
	 * Check out if given node is of red color. Empty node should be considered as BLACK.
	 *
	 * @param node
	 * @return
	 */
    private boolean isRed(Node node) {
        // empty node is black
        if (node == null) {
            return BLACK;
        }

        return node.color;
    }

    /**
	 * Rotate when right red, left black
	 *
	 * @param node
	 * @return
	 */
    private Node leftRotate(Node node) {
        //      node                             nodex
        //     /    \        left rotate        /    \
        //    L    nodex    ------------>    node     XR
        //         /   \                     /    \
        //       XL     XR                  L      XL

        // right child nodex is red, becomes the new root

        Node nodex = node.right;

        // node.right -> XL
        node.right = nodex.left;

        // nodex.left -> node
        nodex.left = node;

        // if node.color is red, nodex must be black?
        // nodex.color should be set to node.color because nodex replace node's position, also
        // need to inherit its color to make the structure stable.
        nodex.color = node.color;

        // node becomes red
        node.color = RED;

        // after rotation, nodex.left = node = red, nodex.left.left = node.left = red

        return nodex;
    }

    /**
	 * Rotate when left red, left.left red
	 *
	 * @param node
	 * @return
	 */
    private Node rightRotate(Node node) {
        //      node                         nodex
        //     /    \      right rotate     /    \
        //    nodex  R    --------->      XL     node
        //    /  \                              /    \
        //   XL   XR                           XR     R

        Node nodex = node.left;

        // node.left -> XR
        node.left = nodex.right;

        // nodex.right -> node
        nodex.right = node;

        // if node.color is red, nodex must be black?
        nodex.color = node.color;

        node.color = RED;

        // after rotation, nodex.left = red, nodex.right = red

        return nodex;
    }

    /**
	 * Insert content to BST
	 *
	 * @param key
	 * @param value
	 */
    public void insert(K key, V value) {
        root = insert(root, key, value);

        // root should be black
        root.color = BLACK;
    }

    /**
	 * Insert content under a specific node and adjust the balance of the tree afterwards, usually called by
	 * <code>{@link #insert(K, V)}</code>.
	 *
	 * @param node
	 * @param key
	 * @param value
	 * @return
	 */
    private Node insert(Node node, K key, V value) {
        if (node == null) {
            size++;

            // The default color of the newly create node is RED
            return new Node(key, value);
        }

        // insert to left branch
        if (key.compareTo(node.key) < 0) {
            node.left = insert(node.left, key, value);
        }
        // insert to right branch
        else if (key.compareTo(node.key) > 0) {
            node.right = insert(node.right, key, value);
        }
        // update the node with given value instead of creating a new one
        else {
            node.value = value;
        }

        // Then adjust node colors to make the tree balanced after insertion.

        return keepBalanced(node);
    }

    /**
	 * Adjust the tree to make it balanced.
	 *
	 * @param node
	 * @return
	 */
    private Node keepBalanced(Node node) {
        // 3 situations

        // Node color being red represents it's not the root of the tree, but neither
        // the empty node.

        // 1. right red, left black

        // If current is the parent of new add node, it should be appended on the right and
        // left branch is null.

        // If the newly inserted node is under node.right, may cause left rotate. Assume current
        // node is a leaf node, after insert the new node, node.right=new_node, node.left=null.

        // If the inserted nodes are in ascending order, it will continuously appending after the
        // right most node which will trigger continuously left rotating.
        if (isRed(node.right) && !isRed(node.left)) {
            node = leftRotate(node);

            // after rotate, node.left=red, node.left.left=black
        }

        // 2. left red, left.left red, children of red node should be black

        // this condition will never gets hit when above condition established and the tree get left rotated.
        // because node.left.left=black

        // If the newly inserted node is under node.left, may cause right rotate. Assume the inserted nodes are
        // in descending order, they will be continuously appending after the left most node which will rigger
        // continuously right rotating.
        if (isRed(node.left) && isRed(node.left.left)) {
            node = rightRotate(node);

            // after rotation, node.left = red, node.right = red
        }

        // 3. both left and right red

        // if the node conducts right rotate, we got a node that its both left and right
        // children are red.
        if (isRed(node.left) && isRed(node.right)) {

            // according to the rule, child node of red node should always be black
            node.color = RED;

            // for nodes of black color, they have no restriction of the color of child nodes.
            node.left.color = BLACK;
            node.right.color = BLACK;
        }

        // So, after a serial of adjustment, the tree will keep balanced finally.

        return node;
    }

    /**
	 * Update the node bind the given key with new value.
	 *
	 * @param key
	 * @param value
	 * @return
	 */
    public void set(K key, V value) {
        Node node = search(root, key);

        if (node == null) {
            throw new IllegalArgumentException(key + " Not Present!");
        }

        node.value = value;
    }

    /**
	 * Search the whole BST for the wanted node binding given key.
	 *
	 * @param key
	 * @return
	 */
    private Node search(K key) {
        return search(root, key);
    }

    /**
	 * Search the specific node for the wanted node binding given key, usually called by
	 * <code>{@link #search(K)}</code>.
	 *
	 * @param node
	 * @param key
	 * @return
	 */
    private Node search(Node node, K key) {
        if (node == null) {
            return null;
        }

        // exactly the wanted node
        if (key.compareTo(node.key) == 0) {
            return node;
        }
        // wanted node on left branch
        else if (key.compareTo(node.key) < 0) {
            return search(node.left, key);
        }

        // wanted node on right branch

        return search(node.right, key);
    }

    /**
	 * Search the whole BST for the node binding the smallest key.
	 *
	 * @param
	 */
    private Node getMin() {
        return getMin(this.root);
    }

    /**
	 * Search under the given node for the node binding the smallest key, usually called by
	 * <code>{@link #getMin()}</code>.
	 *
	 * @param node
	 */
    private Node getMin(Node node) {
        if (node.left == null) {
            return node;
        }

        return getMin(node.left);
    }

    /**
	 * Delete the node binding the given key under the whole BST.
	 *
	 * @param key
	 * @return
	 */
    public V remove(K key) {
        if (isEmpty()) {
            return null;
        }

        Node node = search(root, key);

        if (node != null) {
            root = remove(root, key);

            return node.value;
        }

        return null;
    }

    /**
	 * Delete the node binding the given key under the given node, usually called by
	 * <code>{@link #remove(K)}</code>.
	 *
	 * @param node
	 * @param key
	 * @return
	 */
    private Node remove(Node node, K key) {
        if (node == null) {
            return null;
        }

        // search left
        if (key.compareTo(node.key) < 0) {
            node.left = remove(node.left, key);
        }
        // search right
        else if (key.compareTo(node.key) > 0) {
            node.right = remove(node.right, key);
        } else {
            // exactly the current one is the wanted node

            // 4 situations to be considered separately.

            size--;

            // 1. both empty
            if (node.left == null && node.right == null) {
                return null;
            }

            // 2. left empty
            if (node.left == null) {
                return node.right;
            }

            // 3. right empty
            if (node.right == null) {
                return node.left;
            }

            // 4. both non-empty

            // Find the min node on right branch which should be a node without left branch.
            Node successor = getMin(node.right);

            // Remove the min node from its original position and serves as a successor of current
            // node that going to be removed.
            successor.right = remove(node.right, successor.key);

            successor.left = node.left;

            // unlink the removed node, so that the gc will quickly find it.
            node.left = node.right = null;

            node = successor;
        }

        // Then adjust node colors to make the tree balanced after deletion.

        return keepBalanced(node);
    }

    public boolean contains(K key) {
        return search(root, key) != null;
    }

    public int getSize() {
        return size;
    }


    public boolean isEmpty() {
        return size == 0;
    }

    public int getHeight() {
        return getHeight(root);
    }

    private int getHeight(Node root) {
        if (root == null) {
            return 0;
        }

        int left = getHeight(root.left);
        int right = getHeight(root.right);

        return Math.max(left, right) + 1;
    }

    /**
	 * Verify current tree whether it's a qualified Binary Search Tree or not.
	 *
	 * @return
	 */
    public boolean isBST() {
        ArrayList<K> keys = new ArrayList<>();

        inOrderTraversal(root, keys);

        for (int i = 1; i < keys.size(); i++) {
            if (keys.get(i - 1).compareTo(keys.get(i)) > 0) {
                return false;
            }
        }

        return true;
    }

    private void inOrderTraversal(Node node, ArrayList<K> keys) {
        if (node == null) {
            return;
        }

        inOrderTraversal(node.left, keys);

        keys.add(node.key);

        inOrderTraversal(node.right, keys);
    }

    private class Node {
        K key;
        V value;

        boolean color;

        Node left;
        Node right;

        Node(K key, V value) {
            this.key = key;
            this.value = value;

            this.color = RED;

            this.left = null;
            this.right = null;
        }
    }
}