RBTree.class

RBTree.java

/* RBTree, Java Red Black Tree (left leaning) implementation.
 * Author:
 *       Matthew Clemens
 * Date: 
 *      4/11/13
 * Description: 
 *      Uses the Node class for it's internal nodes. 
 *      Implements the Comparable interface which allows for java Generics.
 *      Exposed are four methods: (traverse, insert, find, remove).
 *      Follows the Red Black tree logic. See file `Algorithm.md` for details.
 * Performance:
 *      O(log n) time for find, insert and remove, and O(n) for print.
 * To Run:
 *       javac RBTree.java; java RBTree
 * Output:
 *       See `main` for program output.
 *
 */
public class RBTree<K extends Comparable<K>, V> {
    
    private Node root = null;
    private boolean red = true;
    private boolean black =false;

    /* Main Program Driver.
     * Each array contains the required values for this assignment.
     * We insert the required values, then print out the Tree and it's state,
     * Then we then remove the required values and print once again. 
    */
    public static void main(String[] args) {

        RBTree<Integer, String> tree = new RBTree<Integer, String>();

        int[] removes =     {9,6,1,3};
        int[] inserts =     {3,2,1,4,5,6,7,
                            16,15,14,13,12,11,10,8,9};
        String[] value =    {"a","b","c","d","e","f","g",
                            "h","i","j","k","l","m","n","o","p"};

        for (int i = 0; i < inserts.length; i++) {
            tree.insert(inserts[i], value[i]);
        }
        tree.traverse();

        for(int i = 0; i <removes.length;i++){
            tree.remove(removes[i]);
        }
        tree.traverse();
    }
    
    // public find, which allows us to correctly start the search 
    // at our tree root
    public V find(K key) {
        return find(root, key);
    }
    
    // recursive find, a standard BST search.
    private V find(Node h, K key) {
        while (h != null) {
            int cmp = key.compareTo(h.key);
            if (cmp == 0) {
                return h.value;
            } else if (cmp < 0) {
                h = h.left;
            } else {
                h = h.right;
            }
        }
        return null;
    }

    // public method which always starts at the root
    public void traverse(){
        postOrderTraverse(root);
        System.out.println();
    }

    // post order traverse our tree
    private void postOrderTraverse(Node p){
        if (p == null) return;
        postOrderTraverse(p.left);
        postOrderTraverse(p.right);
        System.out.println("TRAVERSE [key:"+ p.key+",value:" +p.value +"]");
    }


    
    // public method to insert a generic key and value
    public void insert(K key, V value) {
        root = insert(root, key, value);
        // maintain  our necessary tree invariant
        root.color = black;
    }
    
    // alway connect to its parent with a red link. Thats one 
    // of the reasons for the rotations.
    private Node insert(Node h, K key, V value) {
        // empty tree
        if (h == null) {
            return new Node(key, value);
        }
        
        // If both children are red, flip colors.
        if (isRed(h.left) && isRed(h.right)) {
            colorFlip(h);
        }
        
        int cmp = key.compareTo(h.key);
        if (cmp == 0) {
            h.value = value;
        } else if (cmp < 0) {
            h.left = insert(h.left, key, value);
        } else {
            h.right = insert(h.right, key, value);
        }
        
        // Links are oppisite of correct, flip these nodes
        if (isRed(h.right) && !isRed(h.left)) {
            h = rotateLeft(h);
        }
        
        // If both the left child and its left child are red, rotate right.
        if (isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }
        
        return h;
    }
    
    
    // A need and also useful method which traverses down our 
    // left node links, all the while maintaing our tree invariants by
    private Node removeMin(Node h) {
        // follow the "left" "link" road
        if (h.left == null) {
            return null;
        }
        
        //current node or its left child is red. 
        if (!isRed(h.left) && !isRed(h.left.left)) {
            h = moveRedLeft(h);
        }
        
        // continue down our rabbit hole
        h.left = removeMin(h.left);
        return fixUp(h);
    }
    
    // start on root and work our way down
    public void remove(K key) {
        root = remove(root, key);
        root.color = black;
    }
    
    // perform rotations and color flips on the way down the search path to
    // ensure that the search does not end on a binding red link
    private Node remove(Node h, K key) {
        // bst invariant
        if (key.compareTo(h.key) < 0) {
            // immediate and child left links are wrong need to move 
            if (!isRed(h.left) && !isRed(h.left.left)) {
                h = moveRedLeft(h);
            }
            // continue down the rabbit hole
            h.left = remove(h.left, key);
        } else {
            //rotate left-leaning red links to the right 
            if (isRed(h.left)) {
                h = rotateRight(h);
            }
            
            // found nuthing
            if (key.compareTo(h.key) == 0 && h.right == null) {
                return null;
            }
            
            // our immediate right link is correct, but next level down is 
            // invalid, we have to flip colors... use our helper
            if (!isRed(h.right) && !isRed(h.right.left)) {
                h = moveRedRight(h);
            }
            
            // found but internal node, need to maintain invariants
            if (key.compareTo(h.key) == 0) {
                //replace key and value fields with 
                //minimum node in right subtree
                h.value = find(h.right, min(h.right).key);
                h.key = min(h.right).key;
                h.right = removeMin(h.right);
            } else {
                // continue down the rabbit hole
                h.right = remove(h.right, key);
            }
        }
        
        return fixUp(h);
    }
    
    // maintains our tree invariants for removeMin
    private Node moveRedLeft(Node h) {
        // we need a color flip so that our node stays red
        colorFlip(h);
        
        // check if we have added more reds on accident
        if (isRed(h.right.left)) {
            h.right = rotateRight(h.right);
            h = rotateLeft(h);
            colorFlip(h);
        }

        
        return h;
    }
    
    // maintains our tree invariants
    // not exactly mirror to moveRedLeft, as this is a left leaning tree,
    // and we dont need to double rotate.
    private Node moveRedRight(Node h) {
        colorFlip(h);
        
        if (isRed(h.left.left)) {
            h = rotateRight(h);
            colorFlip(h);
        }
        
        return h;
    }
    
    // We have a right-leaning red link that needs to be rotated to lean to
    // the left. We are simply switching from having the smaller of 
    // the two keys at the root to having the larger of the two keys at 
    // the root.
    private Node rotateLeft(Node h) {
        Node x = h.right;
        h.right = x.left;
        x.left = h;
        
        //preserve its color 
        x.color = h.color;
        h.color = red;
        return x;
    }
    
    // mirror of rotateLeft
    private Node rotateRight(Node h) {
        Node x = h.left;
        h.left = x.right;
        x.right = h;
        
        x.color = h.color;
        h.color = red;
        return x;
    }
    
    // for our delete method
    private Node fixUp(Node h) {
        // no right red links!
        if (isRed(h.right)) {
            h = rotateLeft(h);
        }
        
        // might have added  2 red links in a row
        if ( isRed(h.left) && isRed(h.left.left)) {
            h = rotateRight(h);
        }
        
        // need to "blackify" our links
        if (isRed(h.left) && isRed(h.right)) {
            colorFlip(h);
        }
        // at this point we may be unbalanced. Might have right-leaning red
        // links, or double connected nodes.
        return h;
    }

    /******** Helpers **********/

    private boolean isRed(Node h) {
        return h != null && h.color == red;
    }
    
    private void colorFlip(Node h) {
        // flip colors
        h.color = !h.color;
        h.left.color = !h.left.color;
        h.right.color = !h.right.color;
        // Sanity check, make sure were always black
        assert(h.left.color = black);
        assert(h.right.color = black);
    }
    
    // needed for remove method, as it needs to traverse up and down
    private Node min(Node h) {
        while (h.left != null) {
            h = h.left;
        }
        return h;
    }
    
    
    // Standard Node representation, but agumented with a boolean 
    // used for representing the current color of the node, and also
    // enhanced by using generics!
    private class Node {
        private Node left, right = null;
        private K key;
        private V value;
        private boolean color = red;
        
        public Node(K key, V value) {
            this.key = key;
            this.value = value;
        }
    }
}