/**
 * Red-black tree implementation loosely implementing the ideas in CLRS.
 * 
 * @author Gerth Stølting Brodal
 */

public class RedBlackTree<Element extends Comparable>
{
    /**
     * The root of the red black tree; null if empty tree
     */
    private RedBlackNode<Element> root;

    /**
     * Constructor creating an empty red black tree
     */    
    public RedBlackTree() {
        root = null;
    };


    /**
     * Return if the red black tree is empty
     */    
    public boolean isEmpty() {
        return root == null;
    };

    /**
     * Simple pretty print of a red black tree
     */
    private void print() {
        System.out.println("Red black tree:");
        print_subtree(root, 0);        
    }
        
    /** 
     * Recursive pretty print subtree rooted at x with indentation; red nodes are printed in []
     */
    private void print_subtree(RedBlackNode<Element> x, int indent) {
        if (x != null) {
            print_subtree(x.getRight(), indent + 1);
            for (int i=0; i < indent; i++) System.out.print("   "); 
            if (x.getColor() == Color.RED)
                System.out.println("[" + x.getElement() + "]");
            else    
                System.out.println(x.getElement());
            print_subtree(x.getLeft(), indent + 1);
        }
    };
    
    /**
     * Print inorder traversal of subtree rooted at x
     */
    private void inorderTreeWalk(RedBlackNode<Element> x) {
        if (x != null) {
            inorderTreeWalk(x.getLeft());
            System.out.print(x.getElement() + " ");
            inorderTreeWalk(x.getRight());
        }
    };
    
    /**
     * Print preorder traversal of subtree rooted at x
     */
    private void preorderTreeWalk(RedBlackNode<Element> x) {
        if (x != null) {
            System.out.print(x.getElement() + " ");
            preorderTreeWalk(x.getLeft());
            preorderTreeWalk(x.getRight());
        }
    };

    /**
     * Print postorder traversal of subtree rooted at x
     */
    private void postorderTreeWalk(RedBlackNode<Element> x) {
        if (x != null) {
            postorderTreeWalk(x.getLeft());
            postorderTreeWalk(x.getRight());
            System.out.print(x.getElement() + " ");
        }
    };

    /**
     * Find and return element in red black tree equal to k, otherwise null is returned
     */
    public Element search(Element k) {
        // RedBlackNode<Element> x = recursiveTreeSearch(root, k);
        RedBlackNode<Element> x = iterativeTreeSearch(root, k);
        if (x == null) return null;
        return x.getElement();        
    }
    
    /**
     * Find node in subtree rooted at x with element equal to k, otherwise null is returned
     */
    private RedBlackNode<Element> recursiveTreeSearch(RedBlackNode<Element> x, Element k) {
        if (x == null || x.getElement().compareTo(k) == 0) 
            return x;
        if (x.getElement().compareTo(k) > 0)
            return recursiveTreeSearch(x.getLeft(), k);
        return recursiveTreeSearch(x.getRight(), k);
    }
    
    /**
     * Find node in subtree rooted at x with element equal to k, otherwise null is returned
     */
    private RedBlackNode<Element> iterativeTreeSearch(RedBlackNode<Element> x, Element k) {
        while (x != null && x.getElement().compareTo(k) != 0) 
            if (x.getElement().compareTo(k) > 0)
                x = x.getLeft();
            else 
                x = x.getRight();
        return x;
    }

    /**
     * Find minimum element in red black tree, returns null if empty red black tree
     */
    public Element minimum() {
        if (root == null) return null;
        return treeMinimum(root).getElement();
    }
        
    /**
     * Find node with minimum element in non-empty subtree rooted at node x
     */
    private RedBlackNode<Element> treeMinimum(RedBlackNode<Element> x) {
        while (x.getLeft() != null)
            x = x.getLeft();
        return x;       
    }
    
    /**
     * Find maximum element in red black tree, returns null if empty red black tree
     */
    public Element maximum() {
        if (root == null) return null;
        return treeMaximum(root).getElement();
    }
     
    /**
     * Find node with maximum element in non-empty subtree rooted at node x
     */
    private RedBlackNode<Element> treeMaximum(RedBlackNode<Element> x) {
        while (x.getRight() != null)
            x = x.getRight();
        return x;       
    }
    
    /**
     * Return successor node to node x; 
     * returns null if x is node with maximum element in red black tree
     */    
    private RedBlackNode<Element> treeSuccessor(RedBlackNode<Element> x) {
        if (x.getRight() != null)
            return treeMinimum(x.getRight());
        RedBlackNode<Element> y = x.getParent();
        while (y != null && x == y.getRight()) {
            x = y;
            y = y.getParent();
        }
        return y;
    }
   
    /**
     * Compute depth of red black tree
     */    
    public int depth() {
        return depth(root);
    };
    
    /**
     * Compute depth of subtree rooted at node x
     */    
    private int depth(RedBlackNode<Element> x) {
        if (x == null) {
            return 0;
        }
        return 1 + Math.max(depth(x.getLeft()), depth(x.getRight()));
    };
    
    /**
     * Insert element e into red black tree 
     * (always creates a new node, also if e is already in the tree)
     */    
    public void insert(Element k) {
        RedBlackNode<Element> x = new RedBlackNode<Element>(k);
        if (root == null) {
            root = x;
        } else {
            RedBlackNode<Element> v = root;  // v is current node during top-down search
            while (true) {
                if (v.getElement().compareTo(k) < 0) {
                    if (v.getRight() != null) {
                        v = v.getRight();
                    } else {
                        v.setRight(x);
                        break;
                    }
                } else {
                    if (v.getLeft() != null) {
                        v = v.getLeft();
                    } else {
                        v.setLeft(x);
                        break;
                    }
                } 
            }
        }
        insert_fixup(x);  // fix red-black color invariants
    };
        
    /**
     * Given RED colored node x, fix color invariants up through tree:
     * x is the single node violating color invariants wrt to the edge to parent,
     * i.e. x can be a red root or a red node with red parent 
     * (and x can be red with black parent, then nothing needs to fixed)
     */
    private void insert_fixup(RedBlackNode<Element> x) {
        while (x != root) {
            RedBlackNode<Element> parent = x.getParent();
            if (parent.getColor() == Color.BLACK) return;
            RedBlackNode<Element> 
                grandparent = parent.getParent(),
                sibling = parent == grandparent.getLeft() ? grandparent.getRight() : grandparent.getLeft();
            if (sibling != null && sibling.getColor() == Color.RED) {
                parent.setColor(Color.BLACK);
                sibling.setColor(Color.BLACK);
                grandparent.setColor(Color.RED);
                x = grandparent;
                continue;
            } 
            // x and parent = red, sibling = black/null
            grandparent.setColor(Color.RED);
            if (grandparent.getLeft() == parent) {
                if (x == parent.getLeft()) {
                    // x = grandparent -> left -> left
                    parent.setColor(Color.BLACK);
                } else {
                    // x = grandparent -> left -> right
                    x.setColor(Color.BLACK);
                    rotate_left(parent);
                }
                rotate_right(grandparent);
            } else {
                if (x == parent.getLeft()) {
                    // x = grandparent -> right -> left
                    x.setColor(Color.BLACK);
                    rotate_right(parent);
                } else {
                    // x = grandparent -> right -> right
                    parent.setColor(Color.BLACK);
                }   
                rotate_left(grandparent);
            }
            return;
        }
        root.setColor(Color.BLACK);
    }
    
    /**
     * Rotate right at node x; requires left child to exist
     */
    void rotate_right(RedBlackNode<Element> x) {
        assert x != null && x.getLeft() != null;
        
        RedBlackNode<Element> 
            parent = x.getParent(),
            left = x.getLeft();
    
        if (root == x)
            root = left;
        else {
            if (parent.getLeft() == x)
                parent.setLeft(left);
            else
                parent.setRight(left);
        }
        x.setLeft(left.getRight());
        left.setRight(x);
    }
    
    /**
     * Rotate left at node x; requires right child to exist
     */
    void rotate_left(RedBlackNode<Element> x) {
        RedBlackNode<Element> 
            parent = x.getParent(),
            right = x.getRight();
        
        if (parent == null)
            root = right;
        else {
            if (parent.getRight() == x)
                parent.setRight(right);
            else
                parent.setLeft(right);
        }
        x.setRight(right.getLeft());
        right.setLeft(x);
    }
    
    private enum Direction { LEFT, RIGHT };
    
    /**
     * Delete one node with element k from red black tree
     * (does not change the tree, if k is not in the tree)
     */
    public void delete(Element k) {
        if (root != null) {
            RedBlackNode<Element> z = iterativeTreeSearch(root, k);
            // If k was in the tree at node z 
            if (z != null) {
                // If z has two children, swap with succcesor of z
                if (z.getRight() != null && z.getLeft() != null) {
                    RedBlackNode<Element> y = treeMinimum(z.getRight());
                    z.setElement(y.getElement());
                    z = y;
                }
                // z is node to be deleted with at most one child y, child of parent p in direction d
                RedBlackNode<Element> 
                    p = z.getParent(),
                    y;
                Direction d = p == null || z == p.getLeft() ? Direction.LEFT : Direction.RIGHT;    

                if (z.getRight() == null) {
                    y = z.getLeft();
                    z.setLeft(null);
                } else {
                    y = z.getRight();
                    z.setRight(null);
                }
                
                if (root == z) {
                    root = y;
                    if (root != null)
                        root.setColor(Color.BLACK);
                } else {
                    // make y new child of p
                    if (d == Direction.LEFT) // z was left child of p
                        p.setLeft(y);
                    else // z was right child of p
                        p.setRight(y);
 
                    if (z.getColor() == Color.BLACK)
                        delete_fixup(p, d);
                }
            }
        }
    }

    /**
     * Fix color problem after delete: 
     * Missing black node on edge from v to child in direction d to satisfy black height property (child can be null)
     */
    private void delete_fixup(RedBlackNode<Element> v, Direction d) {
        while (true) {
            if (d == Direction.LEFT) {
                RedBlackNode<Element> 
                    p = v.getParent(), 
                    r = v.getRight(), 
                    rl = r.getLeft(), 
                    rr = r.getRight();
                Color 
                    c = v.getColor(),
                    cr = r.getColor(),
                    crl = rl != null ? rl.getColor() : Color.BLACK,
                    crr = rr != null ? rr.getColor() : Color.BLACK;
                
                if (cr == Color.RED) { // CLRS Case 1
                    v.setColor(Color.RED);
                    r.setColor(Color.BLACK);
                    rotate_left(v);
                } else { // v's sibling is black, CLRS Cases 2-4
                    if (crl == Color.BLACK && crr == Color.BLACK) { // CLRS Case 2
                        r.setColor(Color.RED);
                        if (c == Color.RED) {
                            v.setColor(Color.BLACK);
                            return;
                        } else {
                            if (v == root) return;
                            d = p.getLeft() == v ? Direction.LEFT : Direction.RIGHT;
                            v = p;
                        }
                    } else if (crr == Color.BLACK) { // CLRS Case 3
                        rl.setColor(Color.BLACK);
                        r.setColor(Color.RED);
                        rotate_right(r);
                    } else { // CLRS Case 4
                        v.setColor(Color.BLACK);
                        r.setColor(c);
                        rr.setColor(Color.BLACK);
                        rotate_left(v);
                        return;
                    }
                }
            } else { // symmetric for right side
                RedBlackNode<Element> 
                    p = v.getParent(), 
                    l = v.getLeft(), 
                    lr = l.getRight(), 
                    ll = l.getLeft();
                Color 
                    c = v.getColor(),
                    cl = l.getColor(),
                    clr = lr != null ? lr.getColor() : Color.BLACK,
                    cll = ll != null ? ll.getColor() : Color.BLACK;
                
                if (cl == Color.RED) { // CLRS Case 1
                    v.setColor(Color.RED);
                    l.setColor(Color.BLACK);
                    rotate_right(v);
                } else { // v's sibling is black, CLRS Cases 2-4
                    if (clr == Color.BLACK && cll == Color.BLACK) { // CLRS Case 2
                        l.setColor(Color.RED);
                        if (c == Color.RED) {
                            v.setColor(Color.BLACK);
                            return;
                        } else {
                            if (v == root) return;
                            d = p.getLeft() == v ? Direction.LEFT : Direction.RIGHT;
                            v = p;
                        }
                    } else if (cll == Color.BLACK) { // CLRS Case 3
                        lr.setColor(Color.BLACK);
                        l.setColor(Color.RED);
                        rotate_left(l);
                    } else { // CLRS Case 4
                        v.setColor(Color.BLACK);
                        l.setColor(c);
                        ll.setColor(Color.BLACK);
                        rotate_right(v);
                        return;
                    }
                }
            }
        }
    }
    
    
    /**
     * Validate if the tree is a valid red black tree:
     *   (1) left and right children point to the correct parent,
     *   (2) keys at the nodes appear in inorder, and
     *   (3) validate color invariants
     */
    public void validate() {
        if (root != null) {
            assert root.getParent() == null;
            validate_references(root);        
        }
        validate_inorder(root, null);
        validate_color(root);
    }
    
    /**
     * Traverse subtree rooted at node x and verify all pointers have are set properly
     */
    private void validate_references(RedBlackNode<Element> x) {
        // System.out.println("validating references " + x.getElement());
        if (x == null) return;
        RedBlackNode<Element> 
            left = x.getLeft(),
            right = x.getRight();
        if (left != null) {
            assert left != x;
            assert left.getParent() == x;
            validate_references(left);
        }
        if (right != null) {
            assert right != x;
            assert right.getParent() == x;
            validate_references(right);
        }     
    }
    
    /**
     * Validate search tree satisfies inorder; returns last key verified
     */
    private Element validate_inorder(RedBlackNode<Element> x, Element last) {
        if (x == null) return last;
        last = validate_inorder(x.getLeft(), last);
        Element e = x.getElement();
        assert e != null;
        assert last == null || e.compareTo(last) >= 0;
        return validate_inorder(x.getRight(), e);
    }
    
    /**
     * Validate if red-black tree satsifies red-black invariants; returns black depth
     */
    int validate_color(RedBlackNode<Element> x) {
        if (x == null) return 0;
        int left_depth = validate_color(x.getLeft());
        int right_depth = validate_color(x.getRight());
        assert left_depth == right_depth;
        if (x.getColor() == Color.RED)
            assert x.getParent() != null && x.getParent().getColor() == Color.BLACK;
        return x.getColor() == Color.RED ? left_depth : left_depth + 1;
    }
    
    /**
     * Test all the methods of the class
     */
    public static void test() {
        RedBlackTree<Integer> T = new RedBlackTree<Integer>();

        System.out.println("Insertion increasing integers");
        for (int i=1; i < 25; i++) {
            T.insert(i);
            T.validate();
        }
        T.print();

        System.out.println("Insertion random integers");
        T = new RedBlackTree<Integer>();
        for (int i=0; i<10; i++)
            T.insert((int)(1 + Math.random() * 100));
        T.validate();
        T.print();
        
        System.out.print("Inorder tree walk: ");
        T.inorderTreeWalk(T.root);
        System.out.println();
        
        System.out.print("Preorder tree walk: ");
        T.preorderTreeWalk(T.root);
        System.out.println();
        
        System.out.print("Postorder tree walk: ");
        T.postorderTreeWalk(T.root);
        System.out.println();
        
        System.out.println("Depth = " + T.depth());
        for (int k=40; k<50; k++)
            System.out.println("Search(" + k + "): " + T.search(k));
        
        System.out.println("Minimum: " + T.minimum());
        System.out.println("Maximum: " + T.maximum());
        System.out.println("3rd smallest: " + T.treeSuccessor(T.treeSuccessor(T.treeMinimum(T.root))).getElement());
 
        while (! T.isEmpty()) {
            Integer k = T.minimum();
            System.out.println("Deleting element " + k);
            T.delete(k);
            T.print();
            T.validate();
        }
    
        System.out.print("Stress test...");
        for (int i=0; i<10000; i++)
            T.insert((int)(1 + Math.random() * 1000));
        T.validate();         
        while (! T.isEmpty()) {
            T.delete((int)(1 + Math.random() * 1000));
            T.validate();
        }
        System.out.print("done");
    }
}