数据结构学习之java手写红黑树（左旋、右旋、插入后平衡）

import com.sun.org.apache.regexp.internal.RE;

import java.util.Comparator;

/**
 * @Author PiHao
 * @Date 2020-07-19 18:34
 *
 * 1、创建RbTree,定义颜色
 * 2、创建内部类RbNode
 * 3、辅助方法定义：parentOf(RbNode)、isRed(RbNode)、isBlack(RbNode) setRed(RbNode)、setBlack(RbNode)、inOrderPrint()
 * 4、定义左旋方法：leftRotate(RbNode)
 * 5、定义右旋方法：rightRotate(RbNode)
 * 6、定义外部插入方法：insert(K key,V value)
 * 7、定义内部插入方法：insert(RbNode)
 * 8、插入平衡方法：insertBalance(RbNode)
 * 9、测试红黑树正确性
 *
 */
public class RbTree<K extends Comparator<K>,V> {
    private static final boolean RED = true;
    private static final boolean BLACK = false;

    private RbNode root; //树根的引用

    public RbNode getRoot() {
        return root;
    }

    /**
     * 获取某个节点的父节点
     * @param node
     * @return
     */
    public RbNode parentOf(RbNode node){
        if (node != null){
            return node.parent;
        }
        return null;
    }

    /**
     * 判断节点是否为红色
     * @param node
     * @return
     */
    public boolean isRed(RbNode node){
        if(node != null){
            return node.color == RED;
        }
        return false;
    }

    /**
     * 判断节点是否为黑色
     * @param node
     * @return
     */
    public boolean isBlack(RbNode node){
        if(node != null){
            return node.color == BLACK;
        }
        return false;
    }

    /**
     * 设置节点为红色
     * @param node
     */
    public void setRed(RbNode node){
        if(node != null){
            node.color = RED;
        }
    }

    /**
     * 设置节点为黑色
     * @param node
     */
    public void setBlack(RbNode node){
        if(node != null){
            node.color = BLACK;
        }
    }

    /**
     * 中序打印
     */
    public void inOrderPrint(){
        inOrderPrint(this.root);
    }

    /**
     * 递归打印
     * @param root
     */
    public void inOrderPrint(RbNode root){
        if(root != null){
            inOrderPrint(root.left);
            System.out.println("key: "+ root.key + "value: "+ root.value);
            inOrderPrint(root.right);
        }
    }

    /**
     * 公开的插入方法
     * @param key
     * @param value
     */
    public void insert(K key,V value){
        RbNode node = new RbNode();
        node.setKey(key);
        node.setValue(value);
        node.setColor(true);

        insert(node);
    }

    /**
     * 内部的insert
     * @param node
     */
    public void insert(RbNode node){
        //第一步，查找当前node的父节点
        RbNode parent = null; //定义node的父节点
        RbNode x = this.root;

        while (x != null){
            parent =x;

            //cmp > 0 说明node.key大于x的key,需要到x的右子树查找
            //cmp < 0 说明node.key小于x的key,需要到x的左子树查找
            //cmp = 0 说明node.key等于x的key,替换操作
            int cmp = node.key.compare(node.key, x.key);
            if(cmp > 0){
                x = x.right;
            }else if(cmp < 0){
                x = x.left;
            }else{
                x.setValue(node.getValue());
                return;
            }
        }

        //程序运行到这里，此时的parent就是node的父节点
        node.parent = parent;
        //判断node是parent的左子节点还是右子节点
        if (parent != null){
            int cmp = node.key.compare(node.key, parent.key);
            if(cmp > 0){ //说明node的key大于parent的key,那么就是parent的右子节点
                parent.right = node;
            }else{ //小于0，说明node的key小于parent的key，是parent的左子节点，这里不存在等于0的情况，上面已经替换了
                parent.left = node;
            }
        }else{
            //第一次插入的时候树就是空的
            this.root = node;
        }

        //到这里已经将node插入到了红黑树中，但是这个时候可能破坏了红黑树的平衡性，需要再次将红黑树平衡
        insertBalance(node);

    }

    /**
     * 插入后修复红黑树平衡的方法
     *     |---情景1：红黑树为空树,将根节点染色为黑
     *     |---情景2：插入节点的key已经存在,不用处理,之前已经替换值了
     *     |---情景3：插入节点的父节点为黑色,不用处理,因为没有破坏黑色完美平衡
     *
     *     情景4 需要咱们去处理
     *     |---情景4：插入节点的父节点为红色
     *          |---情景4.1：叔叔节点存在，并且为红色（父-叔 双红）,将爸爸节点和叔叔节点染色为黑,将爷爷节点染色为红,再将爷爷节点设置Wie当前插入节点，进行下一轮处理
     *          |---情景4.2：叔叔节点不存在，或者为黑色，父节点为爷爷节点的左子树
     *               |---情景4.2.1：插入节点为其父节点的左子节点（LL情况） (将爸爸染色为黑,爷爷染色为红,以爷爷节点右旋)
     *               |---情景4.2.2：插入节点为其父节点的右子节点（LR情况）(以爸爸节点左旋回到LL情况,然后指定爸爸节点为当前节点进行下一轮处理)
     *          |---情景4.3：叔叔节点不存在，或者为黑色，父节点为爷爷节点的右子树
     *               |---情景4.3.1：插入节点为其父节点的右子节点（RR情况）(将爸爸节点染色为黑，爷爷染色为红，以爷爷节点左旋)
     *               |---情景4.3.2：插入节点为其父节点的左子节点（RL情况）(以爸爸节点右旋回到RR情况,然后指定爸爸节点为当前节点进行下一轮处理)
     */
    public void insertBalance(RbNode node){
        this.root.setColor(BLACK); //情景一

        RbNode parent = parentOf(node); //爸爸节点
        RbNode grandParent = parentOf(parent);//爷爷节点

        //进入情景四
        if(parent != null && isRed(parent)){
            //如果爸爸节点是红色,那么一定存在爷爷节点,因为爷爷节点不可能是红色，也就是说grandParent不可能为null
            RbNode uncle  = null;
            //找到叔叔节点的位置，left or right
            if(parent == grandParent.left){  //父节点为爷爷节点的左子数树
                uncle = grandParent.right;

                //情景4.1(叔父双红)
                if(uncle != null && isRed(uncle)){
                    setBlack(parent); //将爸爸节点设置为黑色
                    setBlack(uncle); //将叔叔节点设置为黑色
                    setRed(grandParent); //将爷爷节点设置为红色
                    insertBalance(grandParent);//将爷爷节点设置为当前节点，进行下一轮处理
                    return;
                }

                //情景4.2：叔叔节点不存在，或者为黑色，
                if(uncle == null || isBlack(uncle)){

                    //情景4.2.1：插入节点为其父节点的左子节点（LL情况）
                    if(node == parent.left){
                        setBlack(parent); //将爸爸节点染色为黑
                        setRed(grandParent); //将爷爷节点染色为红
                        rightRotate(grandParent); //以爷爷节点右旋
                        return;
                    }

                    //情景4.2.2：插入节点为其父节点的右子节点（LR情况）
                    if(node == parent.right){
                        leftRotate(parent); //先以爸爸节点左旋，得到LL情况
                        insertBalance(parent);//再以爸爸节点为当前节点进行下一轮处理
                        return;
                    }
                }


            }else{ //父节点为爷爷节点的右子树
                uncle = grandParent.left;

                //情景4.1(叔父双红)
                if(uncle != null && isRed(uncle)){
                    setBlack(parent); //将爸爸节点设置为黑色
                    setBlack(uncle); //将叔叔节点设置为黑色
                    setRed(grandParent); //将爷爷节点设置为红色
                    insertBalance(grandParent);//将爷爷节点设置为当前节点，进行下一轮处理
                    return;
                }

                //情景4.3：叔叔节点不存在，或者为黑色
                if(uncle == null || isBlack(uncle)){
                    //情景4.3.1：插入节点为其父节点的右子节点（RR情况）
                    if(node == parent.right){
                        setBlack(parent); //将爸爸染色为黑
                        setRed(grandParent);//将爷爷染色为红
                        leftRotate(grandParent);//将爷爷节点左旋
                        return;
                    }

                    //情景4.3.2：插入节点为其父节点的左子节点（RL情况）
                    if(node == parent.left){
                        rightRotate(parent); //将爸爸节点右旋，得到RR情况，然后以爸爸节点为当前节点进行下一轮操作
                        insertBalance(parent);
                        return;
                    }
                }


            }

        }

    }

    /**
     * 左旋方法
     * 左旋示意图：左旋x节点
     *    p                   p
     *    |                   |
     *    x                   y
     *   / \         ---->   / \
     *  lx  y               x   ry
     *     / \             / \
     *    ly  ry          lx  ly
     *
     * 左旋做了几件事？
     * 1.将y的左子节点赋值给x的右边，并且把x设置为y的左子节点的父节点
     * 2.将x的父节点（非空时）指向y，更新y的父节点为x的父节点
     * 3.将y的左子节点指向x，更新x的父节点为y
     */
    public void leftRotate(RbNode x){
        RbNode y = x.right;
        if(y != null){
            //将y的左子节点复制给x的右子节点，将x复制给y左子节点的父节点
            x.right = y.left;
            y.left.parent = x;
        }

        //将x的父节点（非空时）指向y(x的父节点的左子节点或者右子节点为y)，更新y的父节点为x的父节点,
        if(x.parent != null){
            y.parent = x.parent;
            //判断x是P的左子节点还是右子节点
            if(x == x.parent.left){
                x.parent.left = y;
            }else{
                x.parent.right = y;
            }
        }else{
            //此时说明x就是root,此时需要更新y为根节点，并将y的父节点更新为null;
            this.root = y;
            y.parent = null;
        }

        //将y的左子节点指向x，更新x的父节点为y
        y.left = x;
        x.parent = y;
    }

    /**
     * 右旋方法
     * 右旋示意图：右旋y节点
     *
     *    p                       p
     *    |                       |
     *    y                       x
     *   / \          ---->      / \
     *  x   ry                  lx  y
     * / \                         / \
     *lx  ly                      ly  ry
     *
     * 右旋都做了几件事？
     * 1.将x的右子节点 赋值 给了 y 的左子节点，并且更新x的右子节点的父节点为 y
     * 2.将y的父节点（不为空时）指向x，更新x的父节点为y的父节点
     * 3.将x的右子节点指向y，更新y的父节点为x
     */
    public void rightRotate(RbNode y){
        RbNode x = y.left;
        if(x != null){
            //将x的右子节点 赋值 给了 y 的左子节点，并且更新x的右子节点的父节点为 y
            y.left = x.right;
            x.right.parent = y;
        }
        //将y的父节点（不为空时）指向x，更新x的父节点为y的父节点
        if(y.parent != null){
            x.parent = y.parent;
            if(y.parent.left == y){
                y.parent.left = x;
            }else{
                y.parent.right = x;
            }
        }else{
            //此时说明y就是root,需要跟新x为root节点，并将x的父节点设置为null
            this.root = x;
            x.parent = null;
        }

        //将x的右子节点指向y，更新y的父节点为x
        x.right = y;
        y.parent = x;
    }



    /**
     * 静态内部类，类似hashMap中的Node
     * @param <K>
     * @param <V>
     */
    public static class RbNode<K extends Comparator<K>,V>{
        private RbNode parent;
        private RbNode left;
        private RbNode right;
        private K key;
        private V value;
        private boolean color;

       public  RbNode(){}

        public RbNode(RbNode parent, RbNode left, RbNode right, K key, V value, boolean color) {
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.key = key;
            this.value = value;
            this.color = color;
        }

        public RbNode getParent() {
            return parent;
        }

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

        public RbNode getLeft() {
            return left;
        }

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

        public RbNode getRight() {
            return right;
        }

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

        public K getKey() {
            return key;
        }

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

        public V getValue() {
            return value;
        }

        public void setValue(V value) {
            this.value = value;
        }

        public boolean isColor() {
            return color;
        }

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


}