RBTree

Main.java

package com.gemini.demo;

public class Main {

    public static void main(String[] args) {
	    // write your code here
        TreeNode root = new TreeNode();

        root.setColor(TreeNode.COLOR_BLACK);
        root.setData(50);

        root = root.insert(10);
        root = root.insert(20);
        root = root.insert(30);
        root = root.insert(40);
        root = root.insert(50);
        root = root.insert(60);
        root = root.insert(70);
        root = root.insert(80);
        root = root.insert(90);
        root = root.insert(100);
        System.out.println(root);
    }
}

TreeNode.java

package com.gemini.demo;

public class TreeNode {
    public static final int COLOR_RED = 0;
    public static final int COLOR_BLACK = 1;

    private TreeNode left;
    private TreeNode right;
    private TreeNode parent;

    private int color;
    private int data;

    public TreeNode() {
        this.color = COLOR_RED;
        this.left = null;
        this.right = null;
        this.parent = null;
    }

    public TreeNode getLeft() {
        return left;
    }

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

    public TreeNode getRight() {
        return right;
    }

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

    public int getColor() {
        return color;
    }

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

    public int getData() {
        return data;
    }

    public void setData(int data) {
        this.data = data;
    }

    public TreeNode grandParent() {
        return this.parent.parent;
    }

    public TreeNode uncle() {
        if (this.parent == this.grandParent().left) {
            return this.grandParent().right;
        } else {
            return this.grandParent().left;
        }
    }

    /**
     * 插入一个值，返回 Root
     * @return
     */
    public TreeNode insert(int n) {
        TreeNode child = TreeNode.createNode(this, n);
        if (child != null) {
            insertNode(child);
        }
        return this.findRoot();
    }

    /**
     * 把 node 加入到这棵树中去
     * @param child 子节点
     * @return 根节点
     */
    public void insertNode(TreeNode child) {
        TreeNode parent = child.parent;

        if (parent == null) {
            child.color = COLOR_BLACK;
            return;
        }

        // 如果父亲的节点为黑色
        if (parent.color == COLOR_BLACK) {
            return;
        }

        // 如果父亲节点为红色，然后叔叔节点的颜色也为红色
        if (child.uncle() != null &&
            child.uncle().color == COLOR_RED) {
            child.parent.color = COLOR_BLACK;
            child.uncle().color = COLOR_BLACK;
            child.grandParent().color = COLOR_RED;
            insertNode(child.grandParent());
        } else {
            rotateInsertNode(child);
        }
    }

    private void rotateInsertNode(TreeNode child) {
        TreeNode node = child;
        if (child == child.parent.right &&
            child.parent == child.grandParent().left) {
            node = child.parent.rotateLeft();
            node = node.left;
        } else if (child == child.parent.left &&
            child.parent == child.grandParent().right) {
            node = child.parent.rotateRight();
            node = node.right;
        }
        checkResult(node);
    }

    private TreeNode rotateLeft() {
        TreeNode nextParent = this.right;
        this.right = this.right.left;
        nextParent.left = this;
        if (this.parent != null) {
            if (this.parent.left == this) {
                this.parent.left = nextParent;
            } else {
                this.parent.right = nextParent;
            }
        }
        nextParent.parent = this.parent;
        this.parent = nextParent;
        return nextParent;
    }

    private TreeNode rotateRight() {
        TreeNode nextParent = this.left;
        this.left = nextParent.right;
        nextParent.right = this;
        if (this.parent != null) {
            if (this.parent.right == this) {
                this.parent.right = nextParent;
            } else {
                this.parent.left = nextParent;
            }
        }
        nextParent.parent = this.parent;
        this.parent = nextParent;
        return nextParent;
    }

    /**
     * 检查阶段，看是否需要旋转
     * @param child
     * @return
     */
    private void checkResult(TreeNode child) {
        child.parent.color = COLOR_BLACK;
        child.grandParent().color = COLOR_RED;
        if (child == child.parent.left && child.parent == child.grandParent().left) {
            child.grandParent().rotateRight();
        } else {
            child.grandParent().rotateLeft();
        }
    }

    private static TreeNode createNode(TreeNode root, int n) {
        TreeNode parent = root.findParentForNewValue(n);
        // 待插入的节点已经存在
        if (parent == null) return null;

        // 我们设定插入的孩子必须永远为红色
        TreeNode child = new TreeNode();
        child.color = COLOR_RED;
        child.data = n;
        child.parent = parent;
        if (parent.data > n) {
            parent.left = child;
        } else {
            parent.right = child;
        }
        return child;
    }


    /**
     * 为新插入的值找一个合适的 parent
     * @param n 值
     * @return 新插入的n值的父亲
     */
    private TreeNode findParentForNewValue(int n) {
        if (this.data == n) {
            return null;
        } else if (this.data > n) {
            if (this.left == null) return this;
            return this.left.findParentForNewValue(n);
        } else {
            if (this.right == null) return this;
            return this.right.findParentForNewValue(n);
        }
    }

    private TreeNode findRoot() {
        TreeNode node = this;
        while (node.parent != null) {
           node = node.parent;
        }
        return node;
    }

}