数据结构算法-红黑树

1.红黑树特性

‌1. 节点颜色‌：红黑树的每个节点可以是红色或黑色。根节点必须是黑色，且每个叶子节点（NIL或NULL节点）必须是黑色‌。
2‌. 红色节点的子节点‌：如果一个节点是红色的，则它的子节点必须是黑色的，即不存在连续的红色节点‌。
‌3. 路径长度相等‌：从任一节点到其所有后代叶子节点的路径上，黑色节点的数量必须相同。这意味着没有一条路径会比其他路径长出两倍，从而保持树的平衡‌。
4‌. 根节点颜色‌：根节点必须是黑色的‌

2. 源码

package com.lmh.calculate.tree.redblack;

/**
 * @author liumenghao4
 * @date 2024/12/2
 **/
public class RedBlackTree {

    TreeNode root;

    static class TreeNode {
        TreeNode parent;
        TreeNode left;
        TreeNode right;
        Color color = Color.RED;
        int key;
        Object value;

        TreeNode sibling() {
            if (parent == null) {
                return null;
            }
            if (isLeftChild()) {
                return parent.right;
            } else {
                return parent.left;
            }
        }

        TreeNode uncle() {
            if (parent == null || parent.parent == null) {
                return null;
            }
            return parent.sibling();
        }

        private boolean isLeftChild() {
            return parent != null && parent.left == this;
        }

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

        public TreeNode(int key, Color color, TreeNode left, TreeNode right) {
            this.left = left;
            this.right = right;
            this.color = color;
            this.key = key;
            if (left != null) {
                left.parent = this;
            }
            if (right != null) {
                right.parent = this;
            }
        }
    }

    enum Color {
        RED, BLACK;
    }

    private boolean isRed(TreeNode node) {
        return node != null && node.color == Color.RED;
    }

    private boolean isBlack(TreeNode node) {
        return node == null || node.color == Color.BLACK;
    }

    private void rightRotate(TreeNode node) {
        TreeNode p = node.parent;
        TreeNode l = node.left;
        TreeNode lr = l.right;
        if (l.right != null) {
            lr.parent = node;
        }
        node.left = lr;

        node.parent = l;
        l.right = node;

        if (p == null) {
            root = l;
        } else if (p.left == node) {
            p.left = l;
        } else {
            p.right = l;
        }
        l.parent = p;
    }

    private void leftRotate(TreeNode node) {
        TreeNode p = node.parent;
        TreeNode r = node.right;
        TreeNode rl = r.left;

        if (rl != null) {
            rl.parent = node;
        }
        node.right = rl;

        r.left = node;
        node.parent = r;

        if (p == null) {
            root = r;
//        } else if (node.isLeftChild()) { // TODO 问题原因，复值错误。为什么会导致右边节点全部缺失？
        } else if (p.left == node) {
            p.left = r;
        } else {
            p.right = r;
        }
        r.parent = p;
    }

    public void put(int key, Object value) {
        TreeNode inserted = new TreeNode(key, value);
        if (root == null) {
            root = inserted;
            root.color = Color.BLACK;
            return;
        }
        TreeNode c = root;
        TreeNode p = null;
        while (c != null) {
            p = c;
            if (key < c.key) {
                c = c.left;
            } else if (key > c.key) {
                c = c.right;
            } else {
                // 更新
                c.value = value;
                return;
            }
        }
        if (key < p.key) {
            p.left = inserted;
        } else {
            p.right = inserted;
        }
        inserted.parent = p;
        // 可能会触发红红
        fixRedRed(inserted);
    }

    /**
     * 修复红红相连
     */
    private void fixRedRed(TreeNode c) {
        TreeNode p = c.parent;
        if (c == root) {
            root.color = Color.BLACK;
            return;
        }
        if (isBlack(c) || isBlack(p)) {
            return;
        }
        TreeNode pp = p.parent; // 触发红红, pp必不可能为null
        TreeNode uncle = c.uncle();
        if (isRed(uncle)) {
            p.color = Color.BLACK;
            uncle.color = Color.BLACK;
            pp.color = Color.RED;
            // 祖父可能会触发红红
            fixRedRed(pp);
        } else {
            // 旋转
            if (p.isLeftChild() && c.isLeftChild()) {
                //LL
                rightRotate(pp);
                pp.color = Color.RED;
                p.color = Color.BLACK;
            } else if (p.isLeftChild()) {
                leftRotate(p);
                rightRotate(pp);
                pp.color = Color.RED;
                c.color = Color.BLACK;
            } else if (!c.isLeftChild()) {
                // RR
                leftRotate(pp);
                pp.color = Color.RED;
                p.color = Color.BLACK;
            } else {
                rightRotate(p);
                leftRotate(pp);
                pp.color = Color.RED;
                c.color = Color.BLACK;
            }
        }
    }

    public void remove(int key) {
        TreeNode deleted = find(key);
        if (deleted == null) {
            return;
        }
        doRemove(deleted);
    }

    private void doRemove(TreeNode deleted) {
        // 没有孩子
        TreeNode replaced = findReplaced(deleted);
        TreeNode p = deleted.parent;
        if (replaced == null) {
            if (deleted == root) {
                root = null;
                return;
            }
            if (isBlack(deleted)) {
                // 删除黑色, 少了黑色, 双黑调整, 
                fixDoubleBlack(deleted);
            } else {
                // 删除节点为红色, 直接删除
            }
            if (deleted.isLeftChild()) {
                p.left = null;
            } else {
                p.right = null;
            }
            deleted.parent = null;
            return;
        }
        // 有一个孩子
        if (deleted.left == null || deleted.right == null) {
            if (deleted == root) {
                root.key = replaced.key;
                root.value = replaced.value;
                root.left = root.right = null;
            } else {
                if (deleted.isLeftChild()) {
                    p.left = replaced;
                } else {
                    p.right = replaced;
                }
                replaced.parent = p;
                deleted.left = deleted.right = deleted.parent = null;
                if (isBlack(deleted) && isBlack(replaced)) {
                    // TODO 不会发生双黑情况, 因为双黑的情况下, deleted不会只有一个孩子
                    fixDoubleBlack(deleted);
                } else {
                    replaced.color = Color.BLACK;
                }
            }
            return;
        }
        // 有两个孩子, 与后继节点交换, 并删除后继节点
        Object tv = replaced.value;
        replaced.value = deleted.value; // TODO 个人理解只需要这一行就完事了 
        deleted.value = tv;

        int tk = replaced.key;
        replaced.key = deleted.key;
        deleted.key = tk;

        doRemove(replaced);
    }

    /**
     * 双黑调整
     */
    private void fixDoubleBlack(TreeNode c) {
        if (c == root) {
            return;
        }
        TreeNode sibling = c.sibling();
        TreeNode p = c.parent;
        // 兄弟为红
        if (isRed(sibling)) {
            if (c.isLeftChild()) {
                leftRotate(p);
            } else {
                rightRotate(p);
            }
            p.color = Color.RED;
            sibling.color = Color.BLACK;
            fixDoubleBlack(c);
            return;
        }
        // 兄弟为黑
        if (sibling != null) {
            // 侄子都为黑
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                sibling.color = Color.RED;
                if (isRed(p)) {
                    p.color = Color.BLACK;
                } else {
                    fixDoubleBlack(p);
                }
            } else {
                // 侄子有红
                if (sibling.isLeftChild() && isRed(sibling.left)) {
                    // LL
                    rightRotate(p);
                    sibling.left.color = Color.BLACK;
                    sibling.color = p.color;
                } else if (sibling.isLeftChild() && isRed(sibling.right)) {
                    // LR
                    sibling.right.color = p.color;
                    leftRotate(sibling);
                    rightRotate(p);
                } else if (!sibling.isLeftChild() && isRed(sibling.right)) {
                    // RR
                    sibling.right.color = Color.BLACK;
                    sibling.color = p.color;
                    leftRotate(p);
                } else {
                    // RL
                    sibling.left.color = p.color;
                    rightRotate(sibling);
                    leftRotate(p);
                }
                p.color = Color.BLACK;
            }
        } else {
            // @TODO 实际也不会出现，触发双黑后，兄弟节点不会为 null
            fixDoubleBlack(p);
        }

    }

    /**
     * 找到后继节点
     */
    private TreeNode findReplaced(TreeNode node) {
        if (node.left == null && node.right == null) {
            return null;
        }
        if (node.left == null) {
            return node.right;
        }
        if (node.right == null) {
            return node.left;
        }
        TreeNode p = node.right;
        while (p.left != null) {
            p = p.left;
        }
        return p;
    }

    public boolean contains(int key) {
        return find(key) != null;
    }

    private TreeNode find(int key) {
        TreeNode p = root;
        while (p != null) {
            if (p.key == key) {
                return p;
            } else if (key < p.key) {
                p = p.left;
            } else {
                p = p.right;
            }
        }
        return null;
    }

}