数据结构和算法--红黑树

package com.zwz.avlTree;

import static com.zwz.avlTree.RedBlackTree.Color.BLACK;
import static com.zwz.avlTree.RedBlackTree.Color.RED;

/**
 * 红黑树
 * 自平衡的二叉树
 * 特性：
 * 1。所有节点都有颜色，红色与黑色
 * 2.所有null为黑色
 * 3.红色不相邻--在数中的位置
 * 4。根节点是黑色
 * 5.跟节点到任意叶子节点，黑色节点数量相同
 */
public class RedBlackTree {

    enum Color {
        RED, BLACK;
    }

    private TreeNode root;

    static class TreeNode {
        int key;
        Object value;
        TreeNode left;
        TreeNode right;

        TreeNode parent;
        Color color = RED;


        int height = 1;

        public TreeNode(int key) {
            this.key = key;
        }

        public TreeNode(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        public TreeNode(int key, Object value, TreeNode left, TreeNode right) {
            this.key = key;
            this.value = value;
            this.left = left;
            this.right = right;
        }

        //-------------------工具方法------------------

        /**
         * 判单是否是左孩子
         */
        Boolean isLeftChildren() {
            return parent != null && parent.left == this;
        }

        /**
         * 寻找叔叔节点
         * 父辈节点同级的节点(如果父辈节点是左节点，则返回右)
         */
        TreeNode uncle() {
            if (parent == null || parent.parent == null) {
                return null;
            }
            if (parent.isLeftChildren()) {
                return parent.parent.right;
            } else {
                return parent.parent.left;
            }
        }

        /**
         * 寻找兄弟节点
         */
        TreeNode sibling() {
            if (parent == null) {
                return null;
            }
            if (parent.isLeftChildren()) {
                return parent.right;
            } else {
                return parent.left;
            }
        }

    }

    /**
     * 判断颜色
     *
     * @param node
     * @return
     */
    Boolean isRed(TreeNode node) {
        return node != null && node.color == RED;
    }

    Boolean isBlack(TreeNode node) {
        return node == null || node.color == BLACK;
    }

    /**
     * 右旋
     * 红黑树需要处理parent属性
     *
     * @param pink
     */
    private void turnRight(TreeNode pink) {
        TreeNode parent = pink.parent;
        TreeNode left = pink.left;
        TreeNode leftRigth = left.right;
        //处理parent属性
        if (leftRigth != null) {
            leftRigth.parent = pink;
        }
        left.right = pink;
        //处理parent属性
        left.parent = pink.parent;

        pink.left = leftRigth;
        //处理parent属性
        pink.parent = left;
        //重新处理
        if (parent == null) {
            root = left;
        } else {
            if (parent.left == pink) {
                parent.left = left;
            } else {
                parent.right = left;
            }
        }
    }

    /**
     * 左旋
     * 红黑树需要处理parent属性
     *
     * @param pink
     */
    private void turnLeft(TreeNode pink) {
        TreeNode parent = pink.parent;
        //1.拿到节点的右节点
        TreeNode right = pink.right;
        //2.拿到右节点的左节点
        TreeNode rightLeft = right.left;
        //处理parent属性
        if (rightLeft != null) {
            rightLeft.parent = pink;
        }
        //3.将右节点旋转上来
        //将pink向左旋转下去
        right.left = pink;
        //处理parent属性
        right.parent = pink.parent;


        pink.right = rightLeft;
        //处理parent属性
        pink.parent = rightLeft;
        //重新处理
        if (parent == null) {
            root = right;
        } else {
            if (parent.right == pink) {
                parent.right = right;
            } else {
                parent.left = right;
            }
        }
    }

    /**
     * 节点插入颜色判断(默认红色)
     * 1. 插入节点为根节点 黑
     * 2. 插入节点父节点为黑 红
     * 3. 插入节点父节点为红  触发红红相邻 需要调整平衡
     * 3.1 判断叔叔节点颜色
     * 3.1.1 叔叔节点为红色
     * 1.父节点变黑 叔叔变黑
     * 2.祖父节点变红
     * 3.如果触发红红相邻 递归，一直到根节点
     * 3.1.2 叔叔节点为黑色
     * 判断父节点是左右节点
     * 1.父节点为左节点 插入节点为左节点 父节点变黑 祖父节点变红 祖父节点右旋  LL
     * 2.父节点为左节点 插入节点为右节点 先左旋父节点 插入节点变黑 祖父节点变红 祖父节点右旋 LR
     * 3.父节点为右节点 插入节点为右节点 父节点变黑 祖父节点变红 祖父节点左旋  RR
     * 4.父节点为右节点 插入节点为左节点 先右旋父节点 插入节点变黑 祖父节点变红 祖父节点左旋 RL
     */

    public void put(int key, Object value) {
        TreeNode p = root;
        TreeNode parent = null;
        while (p != null) {
            parent = p;
            if (key > p.key) {
                p = p.right;
            } else if (key < p.key) {
                p = p.left;
            } else {
                p.value = value;
                return;
            }
        }
        TreeNode inserted = new TreeNode(key, value);
        if (parent == null) {
            root = inserted;
        } else {
            if (key > parent.key) {
                parent.left = inserted;
            } else {
                parent.right = inserted;
            }
            inserted.parent = parent;
        }
        fixRedRed(p);
    }

    /**
     * 红红情况调整
     */
    private void fixRedRed(TreeNode node) {
        //1. 根节点
        if (node == root) {
            node.color = BLACK;
            return;
        }
        //2.父节点为黑色
        if (isBlack(node.parent)) {
            return;
        }
        //3.父节点为红色，叔叔节点为红色--变色即可
        TreeNode parent = node.parent; // 父节点
        TreeNode uncle = node.uncle(); //叔叔节点
        TreeNode gParent = parent.parent; //祖父节点
        if (isRed(uncle)) {
            parent.color = BLACK;
            uncle.color = BLACK;
            gParent.color = RED;
            fixRedRed(gParent);
            return;
        }
        //4.父节点为红色，叔叔节点为黑色
        /*
         *          1.父节点为左节点 插入节点为左节点 父节点变黑 祖父节点变红 祖父节点右旋  LL
         *          2.父节点为左节点 插入节点为右节点 先左旋父节点 插入节点变黑 祖父节点变红 祖父节点右旋 LR
         *          3.父节点为右节点 插入节点为右节点 父节点变黑 祖父节点变红 祖父节点左旋  RR
         *          4.父节点为右节点 插入节点为左节点 先右旋父节点 插入节点变黑 祖父节点变红 祖父节点左旋 RL
         */
        //
        if (parent.isLeftChildren()) {
            //1
            if (node.isLeftChildren()) {
                parent.color = BLACK;
                gParent.color = RED;
            } else {
                //2
                turnLeft(parent);
                node.color = BLACK;
                gParent.color = RED;
            }
            turnRight(gParent);
        } else {
            //4
            if (node.isLeftChildren()) {
                turnRight(parent);
                node.color = BLACK;
                gParent.color = RED;
            } else {
                //3
                parent.color = BLACK;
                gParent.color = RED;
            }
            turnLeft(gParent);
        }
    }

    /**
     * 删除节点
     *
     * @param key
     */
    public void remove(int key) {
        TreeNode deleted = find(key);
        if (deleted == null) {
            return;
        }
        doRemove(deleted);

    }

    /**
     * 寻找节点
     *
     * @param key
     * @return
     */
    private TreeNode find(int key) {
        TreeNode p = root;
        while (p != null) {
            if (p.key > key) {
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else {
                return p;
            }
        }
        return null;
    }

    /**
     * 寻找删除后的节点
     */
    private TreeNode findDeleted(TreeNode deleted) {
        if (deleted.left == null && deleted.right == null) {
            return null;
        }
        if (deleted.left == null) {
            return deleted.right;
        }
        if (deleted.right == null) {
            return deleted.left;
        }
        TreeNode s = deleted.right;
        while (s.left != null) {
            s = s.left;
        }
        return s;
    }

    /**
     * 执行删除操作
     *
     * @param deleted
     */
    private void doRemove(TreeNode deleted) {
        TreeNode removed = findDeleted(deleted);
        TreeNode parent = deleted.parent;
        //剩余节点没有子节点
        if (removed == null) {
            //根节点情况
            if (deleted == root) {
                root = null;
            } else {
                //删除的是黑节点，需要左平衡处理
                if (isBlack(deleted)) {
                    //先调整平衡
                    fixBlack(deleted);
                } else {
                    //删除的是红色节点，不需要处理
                }
                //非根节点
                if (deleted.isLeftChildren()) {
                    parent.left = null;
                } else {
                    parent.right = null;
                }
                //致空 有助于垃圾回收
                deleted.parent = null;
            }
            return;
        }
        //剩余节点有一个子节点
        if (deleted.left == null || deleted.right == null) {
            //根节点情况
            if (deleted == root) {
                root.key = removed.key;
                root.value = removed.value;
                root.left = root.right = null;
            } else {
                if (deleted.isLeftChildren()) {
                    parent.left = removed;
                } else {
                    parent.right = removed;
                }
                removed.parent = parent;
                //致空 有助于垃圾回收
                deleted.parent = null;
                deleted.left = null;
                deleted.right = null;

                //删除的是黑色节点，剩下的也是黑色节点，需要做特殊处理
                if (isBlack(deleted) && isBlack(removed)) {
                    //后调整平衡
                    fixBlack(removed);
                } else {
                    //删除的是黑色节点，剩下的红色，只需要变换颜色
                    removed.color = BLACK;
                }
            }
            return;
        }
        //剩余节点有两个子节点，进行递归键值替换，转换成以上两种情况
        int k = deleted.key;
        deleted.key = removed.key;
        removed.key = k;

        Object v = deleted.value;
        deleted.value = removed.value;
        removed.value = v;
        doRemove(removed);

    }

    /**
     * 处理双黑色删除平衡
     *
     * @param deleted
     */
    private void fixBlack(TreeNode deleted) {
        if (deleted == root) {
            return;
        }
        TreeNode parent = deleted.parent;//找到父节点
        TreeNode sibling = deleted.sibling();//找到兄弟节点
        //如果兄弟节点是红色,先进行一次旋转
        if (isRed(sibling)) {
            //此处目的是兄弟颜色置换为黑色
            //左子树 父节点左旋 右子树 父节点右旋
            if (deleted.isLeftChildren()) {
                turnLeft(parent);
            } else {
                turnRight(parent);
            }
            //换颜色
            parent.color = RED;
            sibling.color = BLACK;
            //递归调用，
            fixBlack(deleted);
            return;
        }
        //兄弟节点是黑色，判断兄弟节点左右颜色
        if(sibling != null){
            //判断兄弟节点左右颜色都是黑色
            if(isBlack(sibling.left) && isBlack(sibling.right)){
                sibling.color = RED;
                if(isRed(parent)){
                    parent.color = BLACK;
                }else{
                    fixBlack(parent);
                }
            }else{
                //判断兄弟节点左右颜色至少有一个是红色
                //分为四种情况 LL LR RR RL
                //兄弟是左 兄弟左侧树 LL
                if(sibling.isLeftChildren() && isRed(sibling.left)){
                    //进行一次右旋
                    turnRight(parent);
                    //变色
                    sibling.left.color = BLACK;
                    sibling.color = parent.color;
                }else if(sibling.isLeftChildren() && isRed(sibling.right)){
                    //兄弟是左 兄弟右侧树 LR
                    sibling.right.color = parent.color;
                    //兄弟先左旋
                    turnLeft(sibling);
                    //父级右旋
                    turnRight(parent);
                }else if(!sibling.isLeftChildren() && isRed(sibling.left)){
                    //兄弟是右 兄弟左侧树 RL
                    sibling.left.color = parent.color;
                    turnRight(sibling);
                    turnLeft(parent);
                }else{
                    //兄弟是右 兄弟右侧树 RR
                    //进行一次右旋
                    turnLeft(parent);
                    //变色
                    sibling.right.color = BLACK;
                    sibling.color = parent.color;
                }
                parent.color = BLACK;
            }
        }else{
            fixBlack(parent);
        }
    }
}