红黑树插入/删除 详细流程图和代码实现

插入流程图

删除流程图

 

红黑树代码实现

/**
 * 红黑树
 *
 * @author: zyp
 * @since: 2021/11/25 13:32
 */
public class RedBlackTreeMap<K extends Comparable<K>, V> {
    //    NULL用来标记NULL节点
    private static final Object NULL = new Object();
    private RedBlackTreeNode<K, V> root;

    public V get(K key) {
        if (root == null) {
            return null;
        }
        RedBlackTreeNode<K, V> target = root.find(key);
        return target == null ? null : target.value;
    }

    public V remove(K key) {
        if (root == null) {
            return null;
        }
        RedBlackTreeNode<K, V> target = root.find(key);
        if (target == null) {
            return null;
        }
        V deletedValue = target.value;
//        搜索后继节点, 并将后继节点的值赋值给自己
        target = target.findAppendNodeForDeleteAndCopyValue();
//        现在target是真正要删除的节点, target最多只有1个孩子
        Direction direction = target.getDirection();
        RedBlackTreeNode<K, V> tp = target.parent;
//        把target从树种删除
        RedBlackTreeNode<K, V> replaceNode = target.deleteFromTree();
        if (target.red) {
//            删除的是红色节点, 结束
            return deletedValue;
        }
//        删除的是黑色节点
        if (replaceNode != null) {
//            有顶替节点, 顶替节点一定是红色, 把顶替节点改成黑色
            replaceNode.setBlack();
            return deletedValue;
        }
//        顶替节点(null)成为双黑节点
        handleDoubleBlackNode(tp, null, direction);
        return deletedValue;
    }

    public V put(K key, V value) {
        if (root == null) {
            this.root = new RedBlackTreeNode(false, key, value, null, null, null);
            return null;
        }
        RedBlackTreeNode<K, V> c = root.insert(key);
        if (!c.isNullNode()) {
//            是替换
            return c.replaceValue(value);
        }
//        是新增
        c.replaceValue(value);
//        平衡
        c.balanceInsert();
        return null;
    }


    public int size() {
        return root == null ? 0 : root.size();
    }

    public void clear() {
        root = null;
    }

    /**
     * @param p         p 父节点
     * @param c         c 当前节点, 可null
     * @param direction left 当前节点是父节点的左孩子吗
     * @return void
     * @author zhengyongpan
     * @since 2021/11/26 18:28
     */
    private void handleDoubleBlackNode(RedBlackTreeNode<K, V> p, RedBlackTreeNode<K, V> c, Direction direction) {
        if (c == root) {
//            double black node is root. make it single black.
            c.red = false;
            return;
        }
        RedBlackTreeNode<K, V> siblingNode = direction.isLeft() ? p.right : p.left;
        if (siblingNode.red) {
//                double black node has red sibling. retate tree to make sibling black (roteta parent node)
            p.setRed();
            siblingNode.setBlack();
            p.rotate(direction);
            handleDoubleBlackNode(p, c, direction);
        } else {
            if (siblingNode.isLeftChildBlack() && siblingNode.isRightildBlack()) {
//                  Double black node has blak sibling. 2 black nephews. Push up black level;
                siblingNode.setRed();
                if (p.isRed()) {
                    p.setBlack();
                } else {
//                   Pushing up black level created another double black node.  Repeating ...
                    handleDoubleBlackNode(p.parent, p, p.getDirection());
                }
            } else if (siblingNode.isChildBlack(direction.opposite())) {
//        Double black node has black sibling, but double black node is a left child, and the right nephew is black.  Rotate tree to make opposite nephew red ...
//        Double black node has black sibling, but double black node is a right child, and the left nephew is black.  Rotate tree to make opposite nephew red ...
                siblingNode.setChildBlack(direction);
                siblingNode.setRed();
                siblingNode.rotate(direction.opposite());
                handleDoubleBlackNode(p, c, direction);
            } else {
//        Double black node has black sibling, is a left child, and its right nephew is red. One rotation can fix double-blackness.
//        Double black node has black sibling, is a right child, and its left nephew is red. One rotation can fix double-blackness.
                if (p.isRed()) {
//                        反向外甥节点是红色, 兄弟节点是黑色, 父节点是红色
                    p.setBlack();
                    siblingNode.setRed();
                } else {
//                        反向外甥节点是红色, 兄弟节点是黑色, 父节点是黑色
                }
                siblingNode.setChildBlack(direction.opposite());
                p.rotate(direction);
                root.setBlack();
            }
        }
    }

    /**
     * 验证当前树是否是红黑树
     *
     * @param
     * @return boolean
     * @author zhengyongpan
     * @since 2021/11/26 16:08
     */
    public void validate() {
        if (root == null) {
            return;
        }
        if (root.red) {
            throw new RuntimeException("红色根节点");
        }
        root.validate();
    }


    /**
     * 红黑树节点
     *
     * @author zhengyongpan
     * @since 2021/11/27 14:25
     */
    private class RedBlackTreeNode<K, V> {
        private RedBlackTreeNode<K, V> left;
        private RedBlackTreeNode<K, V> right;
        private RedBlackTreeNode<K, V> parent;
        private K key;
        private boolean red;
        private V value;

        public RedBlackTreeNode(boolean red, K key, V value, RedBlackTreeNode<K, V> parent, RedBlackTreeNode<K, V> left, RedBlackTreeNode<K, V> right) {
            this.left = left;
            this.right = right;
            this.parent = parent;
            this.key = key;
            this.value = value;
            this.red = red;
        }


        /**
         * 搜索目标节点
         *
         * @param key key
         * @return test.RedBlackTreeMap.RedBlackTreeNode<K, V>
         * @author zhengyongpan
         * @since 2021/11/25 18:02
         */
        private RedBlackTreeNode<K, V> find(K key) {
            Direction direction = RedBlackTreeMap.getDirection(this.key, key);
            if (direction == null) {
                return this;
            }
            RedBlackTreeNode<K, V> child = this.getChild(direction);
            if (child == null) {
                return null;
            }
            return child.find(key);
        }

        public V replaceValue(V value) {
            V oldValue = this.value;
            this.value = value;
            return oldValue;
        }

        /**
         * 是否是NULL节点
         *
         * @param
         * @return boolean
         * @author zhengyongpan
         * @since 2021/11/27 14:25
         */
        public boolean isNullNode() {
            return value == NULL;
        }

        /**
         * 插入key
         * 如果树种已经有该key了, 返回已存在的节点
         * 如果树种没有该key, 创建新节点, 插入到树上, 新节点的值为NULL
         *
         * @param key key
         * @return test.RedBlackTreeMap<K, V>.RedBlackTreeNode<K,V>
         * @author zhengyongpan
         * @since 2021/11/27 12:13
         */
        private RedBlackTreeNode<K, V> insert(K key) {
            Direction direction = RedBlackTreeMap.getDirection(this.key, key);
            if (direction == null) {
                return this;
            } else {
                RedBlackTreeNode<K, V> child = this.getChild(direction);
                if (child == null) {
                    RedBlackTreeNode isnertNode = new RedBlackTreeNode(true, key, NULL, null, null, null);
                    this.setChild(direction, isnertNode);
                    return isnertNode;
                }
                return child.insert(key);
            }
        }

        /**
         * 获取真正要删除的后继节点, 把后继节点的值赋值给自己
         * 如果没有子节点, 就返回自己
         *
         * @param
         * @return test.RedBlackTreeMap<K, V>.RedBlackTreeNode<K,V>
         * @author zhengyongpan
         * @since 2021/11/27 11:00
         */
        public RedBlackTreeNode<K, V> findAppendNodeForDeleteAndCopyValue() {
            RedBlackTreeNode<K, V> appendNdoe;
            Direction direction = null;
            if (left != null) {
                direction = Direction.left;
            } else if (right != null) {
                direction = Direction.right;
            }
            if (direction == null) {
                appendNdoe = this;
            } else {
                RedBlackTreeNode<K, V> c = this.getChild(direction);
                direction = direction.opposite();
                while (c.getChild(direction) != null) {
                    c = c.getChild(direction);
                }
                appendNdoe = c;
            }
            this.key = appendNdoe.key;
            this.value = appendNdoe.value;
            return appendNdoe;
        }

        /**
         * 设置孩子
         *
         * @param direction direction
         * @param child     child
         * @return test.RedBlackTreeMap.RedBlackTreeNode<K, V> 被替换的孩子
         * @author zhengyongpan
         * @since 2021/11/27 10:20
         */
        public RedBlackTreeNode<K, V> setChild(Direction direction, RedBlackTreeNode<K, V> child) {
            RedBlackTreeNode<K, V> replaced;
            if (direction.isLeft()) {
                replaced = left;
                left = child;
            } else {
                replaced = right;
                right = child;
            }
            if (child != null) {
                child.parent = this;
            }
            return replaced;
        }

        public RedBlackTreeNode<K, V> getChild(Direction direction) {
            return direction.isLeft() ? left : right;
        }

        public void setBlack() {
            red = false;
        }

        public void setRed() {
            red = true;
        }

        public void setChildBlack(Direction direction) {
            if (direction.isLeft()) {
                if (left != null) {
                    left.setBlack();
                }
            } else {
                if (right != null) {
                    right.setBlack();
                }
            }
        }

        public void setChildRed(Direction direction) {
            if (direction.isLeft()) {
                if (left != null) {
                    left.setRed();
                }
            } else {
                if (right != null) {
                    right.setRed();
                }
            }
        }


        public boolean isChildBlack(Direction direction) {
            return direction.isLeft() ? isLeftChildBlack() : isRightildBlack();
        }

        /**
         * 获取当前节点是父节点的左孩子还是右孩子
         * 如果没有父节点, 返回null
         *
         * @param
         * @return test.RedBlackTreeMap.Direction
         * @author zhengyongpan
         * @since 2021/11/27 12:06
         */
        public Direction getDirection() {
            if (parent == null) {
                return null;
            }
            return parent.left == this ? Direction.left : Direction.right;
        }

        public boolean isBlack() {
            return !red;
        }

        public boolean isRed() {
            return red;
        }

        public boolean isLeftChildBlack() {
            return left == null || left.isBlack();
        }

        public boolean isRightildBlack() {
            return right == null || right.isBlack();
        }

        /**
         * 获取兄弟节点
         *
         * @param
         * @return test.RedBlackTreeMap.RedBlackTreeNode<K, V>
         * @author zhengyongpan
         * @since 2021/11/26 14:15
         */
        public RedBlackTreeNode<K, V> getSiblingNode() {
            if (parent == null) {
                return null;
            }
            if (parent.left == this) {
                return parent.right;
            } else {
                return parent.left;
            }
        }

        /**
         * 把当前节点从树种删除
         * 返回替换自己的节点
         *
         * @param
         * @return void
         * @author zhengyongpan
         * @since 2021/11/26 14:04
         */
        public RedBlackTreeNode<K, V> deleteFromTree() {
            if (left != null && right != null) {
                throw new RuntimeException("both child not null");
            }
            RedBlackTreeNode<K, V> replaceNode;
            if (left == null) {
                replaceNode = right;
            } else {
                replaceNode = left;
            }
            if (replaceNode != null) {
                replaceNode.parent = parent;
            }
            if (parent != null) {
                if (parent.left == this) {
                    parent.left = replaceNode;
                } else {
                    parent.right = replaceNode;
                }
                this.parent = null;
            }
            this.left = null;
            this.right = null;
            return replaceNode;
        }


        @Override
        public String toString() {
            return String.format("c=%s, k=%s, v=%s", red ? "红" : "黑", key, value);
        }

        public int size() {
            return 1 + (left == null ? 0 : left.size()) + (right == null ? 0 : right.size());
        }

        /**
         * 校验是否满足红黑树限制条件
         *
         * @param
         * @return void
         * @author zhengyongpan
         * @since 2021/11/27 12:00
         */
        public void validate() {
            if (red) {
                if (left != null && left.red) {
                    throw new RuntimeException("相邻红色节点");
                }
                if (right != null && right.red) {
                    throw new RuntimeException("相邻红色节点");
                }
            }
            if (left != null) {
                left.validate();
            }
            if (right != null) {
                right.validate();
            }
            if (getBlackDepth(Direction.left) != getBlackDepth(Direction.right)) {
                throw new RuntimeException("左右分支黑色节点数不相同");
            }
        }

        /**
         * 返回子树中黑色节点深度
         *
         * @param direction direction
         * @return int
         * @author zhengyongpan
         * @since 2021/11/27 12:00
         */
        public int getBlackDepth(Direction direction) {
            int depth = 1;
            RedBlackTreeNode<K, V> c = this;
            while (c != null) {
                if (c.isBlack()) {
                    depth++;
                }
                c = c.getChild(direction);
            }
            return depth;
        }

        public void balanceInsert() {
            if (parent == null) {
//                当前节点是根节点, 设置为黑色
                red = false;
                return;
            }
            if (parent.isBlack()) {
//            父节点是黑色
                return;
            }
//          父节点是红色
            RedBlackTreeNode<K, V> pp = parent.parent;
            if (pp == null) {
//                祖父节点不存在, 说明父节点是根节点, 设置为黑色
                parent.red = false;
                return;
            }
            RedBlackTreeNode<K, V> uncleNode = parent.getSiblingNode();
            if (uncleNode != null && uncleNode.red) {
//                叔叔节点是红色, 祖父节点是黑色, 父节点是红色, 将父节点/祖父节点/叔叔节点变色, 将祖父节点视为当前节点, 递归
                parent.setBlack();
                pp.setRed();
                uncleNode.setBlack();
                pp.balanceInsert();
                return;
            }
            Direction direction = getDirection();
            Direction directionP = parent.getDirection();
            if (directionP != direction) {
//                父节点是祖父节点的左孩子, 当前节点是父节点的右孩子, 父节点左旋
//                or 父节点是祖父节点的右孩子, 当前节点是父节点的左孩子, 父节点右旋
                this.setBlack();
                parent.rotate(directionP);
            } else {
                parent.setBlack();
            }
            pp.setRed();
            pp.rotate(directionP.opposite());
        }


        private RedBlackTreeNode<K, V> rotate(Direction direction) {
            if (direction.isLeft()) {
                return rotateLeft();
            } else {
                return rotateRight();
            }
        }

        /**
         * 左旋, 返回新的支点
         *
         * @param root root
         * @param c    c
         * @return test.RedBlackTreeMap.RedBlackTreeNode<K, V>
         * @author zhengyongpan
         * @since 2021/11/25 18:01
         */
        private RedBlackTreeNode<K, V> rotateLeft() {
            RedBlackTreeNode<K, V> c = this;
            if (c.right == null) {
                throw new RuntimeException();
            }
            RedBlackTreeNode<K, V> cr = c.right;
            RedBlackTreeNode<K, V> crl = cr.left;
            RedBlackTreeNode<K, V> cp = c.parent;

            if (cp != null) {
                if (cp.left == c) {
                    cp.left = cr;
                } else {
                    cp.right = cr;
                }
                cr.parent = cp;
            } else {
                cr.parent = null;
                root = (RedBlackTreeNode) cr;
            }

            cr.left = c;
            c.parent = cr;

            c.right = crl;
            if (crl != null) {
                crl.parent = c;
            }
            return cr;
        }

        /**
         * 右旋, 返回新的支点
         *
         * @param root root
         * @param c    c
         * @return test.RedBlackTreeMap.RedBlackTreeNode<K, V>
         * @author zhengyongpan
         * @since 2021/11/25 18:01
         */
        private RedBlackTreeNode<K, V> rotateRight() {
            RedBlackTreeNode<K, V> c = this;
            if (c.left == null) {
                throw new RuntimeException();
            }
            RedBlackTreeNode<K, V> cl = c.left;
            RedBlackTreeNode<K, V> clr = cl.right;
            RedBlackTreeNode<K, V> cp = c.parent;

            if (cp != null) {
                if (cp.left == c) {
                    cp.left = cl;
                } else {
                    cp.right = cl;
                }
                cl.parent = cp;
            } else {
                cl.parent = null;
                root = (RedBlackTreeNode) cl;
            }

            cl.right = c;
            c.parent = cl;

            c.left = clr;
            if (clr != null) {
                clr.parent = c;
            }
            return cl;
        }
    }

    /**
     * 判断key2应该放在key1的左边还是右边
     * null is key1 equals key2
     *
     * @param key1 key1
     * @param key2 key2
     * @return test.RedBlackTreeMap.Direction
     * @author zhengyongpan
     * @since 2021/11/27 10:15
     */
    private static Direction getDirection(Object key1, Object key2) {
        int compare;
        if (key1 == null) {
            if (key2 == null) {
                compare = 0;
            } else {
                compare = -1;
            }
        } else if (key2 == null) {
            compare = 1;
        } else {
            compare = ((Comparable) key1).compareTo(key2);
            if (key1.getClass() != key2.getClass()) {
                throw new RuntimeException(String.format("key type not match: key1=%s, key2=%s", key1.getClass(), key2.getClass()));
            }
        }
        if (compare == 0) {
            return null;
        }
        if (compare < 0) {
            return Direction.right;
        } else {
            return Direction.left;
        }
    }

    @Override
    public String toString() {
        StringJoiner sb = new StringJoiner(", ", "{", "}");
        toString(sb, root);
        return sb.toString();
    }

    private void toString(StringJoiner sb, RedBlackTreeNode node) {
        if (node == null) {
            return;
        }
        toString(sb, node.left);
        sb.add(String.format("%s=%s", node.key, node.value));
        toString(sb, node.right);
    }

    /**
     * 左右方向
     *
     * @author zhengyongpan
     * @since 2021/11/27 14:25
     */
    private static enum Direction {
        left,
        right;
        private static Direction[] directions = new Direction[]{left, right};

        public Direction opposite() {
            return directions[1 - this.ordinal()];
        }

        public boolean isLeft() {
            return this == left;
        }

        public boolean isRight() {
            return this == right;
        }
    }

}