Java 红黑树

红黑树比较复杂，深入了解之后，感觉和魔方一样，都是遇到某种情况下，套用固定的公式来完成即可。一般的实现代码四百行左右，此次合并了某些情况，把代码压缩到了210行左右。

import lombok.Data;

/**
 * @author wzgl
 * @version 1.0
 * @date 2023/5/30 11:49
 */
public class MyRBTree {

    public static void main(String[] args) {
        MyRBTree tree = new MyRBTree();
        int[] arr = {65, 99, 67, 100, 68, 7, 9, 75, 11, 43, 45, 78, 46, 47, 21, 87, 24, 25, 26, 95};
        for (int i : arr) {
            tree.add(i);
            tree.check();
        }
        for (int i : arr) {
            tree.remove(i);
            tree.check();
        }
    }

    /**
     * 测试是否时红黑树
     */
    public void check() {
        CheckRBTree.PrintNodeInfo<Node> info = new CheckRBTree.PrintNodeInfo<>(this.root, MyRBTree.Node::getVal, MyRBTree.Node::getLeft, MyRBTree.Node::getRight);
        info.setRb(Node::isBlack);
        if (!CheckRBTree.checkRBTree(info)) {
            throw new RuntimeException("违反了红黑树的结构");
        }
    }

    private Node root;

    /**
     * 新增
     *
     * @param v 元素
     */
    public void add(int v) {
        Node node = new Node(v);
        if (root == null) {
            root = node;
            root.black = true;
            return;
        }
        Node parent = findParent(v, root);
        node.parent = parent;
        if (v < parent.val) {
            parent.left = node;
        } else {
            parent.right = node;
        }
        if (parent.black) {
            return;
        }
        addFix(node);
    }
    
    /**
     * 删除
     *
     * @param v 元素
     */
    public boolean remove(int v) {
        Node node = find(v, root);
        if (node == null) {
            return false;
        }
        Node del = node;
        if (node.right != null) {
            del = findMin(node.right);
            node.val = del.val;
        }
        removeFix(del);
        del(del);
        return true;
    }

    private void removeFix(Node node) {
        if (!node.black || node == root) {
            node.black = true;
            return;
        }
        Node parent = node.parent;
        boolean itemType = parent.left == node;
        Node brother = itemType ? parent.right : parent.left;
        if (!brother.black) {
            (itemType ? leftRote(brother) : rightRote(brother)).black = true;
            parent.black = false;
            parent = node.parent;
            brother = itemType ? parent.right : parent.left;
        }
        Node bl = brother.left;
        Node br = brother.right;
        if ((bl == null || bl.isBlack()) && (br == null || br.isBlack())) {
            brother.black = false;
            removeFix(parent);
            return;
        }
        if (itemType ^ (bl == null || bl.black)) {
            brother.black = false;
            (itemType ? rightRote(bl) : leftRote(br)).black = true;
            brother = itemType ? parent.right : parent.left;
        }
        (itemType ? leftRote(brother) : rightRote(brother)).black = parent.black;
        parent.black = true;
        (itemType ? brother.right : brother.left).black = true;
    }

    private void del(Node node) {
        Node next = node.left == null ? node.right : node.left;
        Node parent = node.parent;
        if (parent == null) {
            root = next;
            if (root != null) {
                root.black = true;
            }
        } else {
            if (parent.left == node) {
                parent.left = next;
            } else {
                parent.right = next;
            }
        }
        if (next != null) {
            next.parent = parent;
        }
        node.left = null;
        node.right = null;
        node.parent = null;
    }

    private Node findMin(Node node) {
        if (node.left == null) {
            return node;
        }
        return findMin(node.left);
    }

    private Node find(int v, Node node) {
        if (node == null || v == node.val) {
            return node;
        }
        return v < node.val ? find(v, node.left) : find(v, node.right);
    }

    private void addFix(Node node) {
        Node parent = node.parent;
        if (parent == null) {
            node.black = true;
            return;
        }
        if (parent.black) {
            return;
        }
        Node grand = parent.parent;
        boolean parentType = grand.left == parent;
        Node uncle = parentType ? grand.right : grand.left;
        if (uncle != null && !uncle.black) {
            parent.black = true;
            uncle.black = true;
            grand.black = false;
            addFix(grand);
            return;
        }
        node = parentType ^ (parent.left == node) ? (parentType ? leftRote(node) : rightRote(node)) : parent;
        (parentType ? rightRote(node) : leftRote(node)).black = true;
        grand.black = false;
    }

    private Node leftRote(Node node) {
        Node parent = node.parent;
        Node grand = parent.parent;
        Node left = node.left;
        node.left = parent;
        parent.parent = node;
        parent.right = left;
        if (left != null) {
            left.parent = parent;
        }
        node.parent = grand;
        if (grand == null) {
            root = node;
        } else {
            if (grand.left == parent) {
                grand.left = node;
            } else {
                grand.right = node;
            }
        }
        return node;
    }

    private Node rightRote(Node node) {
        Node parent = node.parent;
        Node grand = parent.parent;
        Node right = node.right;
        node.right = parent;
        parent.parent = node;
        parent.left = right;
        if (right != null) {
            right.parent = parent;
        }
        node.parent = grand;
        if (grand == null) {
            root = node;
        } else {
            if (grand.left == parent) {
                grand.left = node;
            } else {
                grand.right = node;
            }
        }
        return node;
    }

    private Node findParent(int v, Node node) {
        Node next = v < node.val ? node.left : node.right;
        if (next == null) {
            return node;
        }
        return findParent(v, next);
    }

    @Data
    public class Node {
        private int val;
        private Node left;
        private Node right;
        private Node parent;
        private boolean black = false;

        public Node(int val) {
            this.val = val;
        }
    }
}

为了校验是否符合红黑树的性质，再写一个check类。来校验是否是红黑树。

import lombok.Data;

import java.util.function.Function;

/**
 * @author wzgl
 * @version 1.0
 * @date 2022/6/2 13:48
 */
public class CheckRBTree {
    
    public static <T> boolean checkRBTree(PrintNodeInfo<T> info) {
        if (info == null) {
            return true;
        }
        T node = info.getNode();
        if (node == null) {
            System.out.println("实际深度 : " + 0 + " ; 理想深度 : " + 0 + "; 总节点树 : " + 0);
            return true;
        }
        int depth = findDepth(node, info);
        int nodeCount = findNodeCount(node, info);
        int t = 0;
        int e = nodeCount;
        while (e != 0) {
            t++;
            e >>= 1;
        }
        System.out.println("实际深度 : " + depth + " ; 理想深度 : " + t + "; 总节点树 : " + nodeCount);
        if (!info.rb.apply(node)) {
            System.out.println("根节点必须为黑色节点");
            return false;
        }
        return checkRBTree(info.getNode(), info) > -1;
    }

    private static <T> int checkRBTree(T node, PrintNodeInfo<T> info) {
        if (node == null) {
            return 0;
        }
        if (!info.rb.apply(node)) {
            T left = info.left.apply(node);
            if (left != null && !info.rb.apply(left)) {
                System.out.println("红节点的左节点必须为黑节点");
                return -1;
            }
            if (left != null && info.val.apply(left) > info.val.apply(node)) {
                System.out.println("左节点必须小于当前节点");
                return -1;
            }
            T right = info.right.apply(node);
            if (right != null && !info.rb.apply(right)) {
                System.out.println("红节点的右节点必须为黑节点");
                return -1;
            }
            if (right != null && info.val.apply(right) < info.val.apply(node)) {
                System.out.println("左节点必须大于当前节点");
                return -1;
            }
        }
        int a = checkRBTree(info.left.apply(node), info);
        int b = checkRBTree(info.right.apply(node), info);
        if (b == -1 || a != b) {
            return -1;
        }
        return a + (info.rb.apply(node) ? 1 : 0);
    }

    private static <T> int findDepth(T node, PrintNodeInfo<T> info) {
        if (node == null) {
            return 0;
        }
        return Math.max(findDepth(info.left.apply(node), info), findDepth(info.right.apply(node), info)) + 1;
    }

    private static <T> int findNodeCount(T node, PrintNodeInfo<T> info) {
        if (node == null) {
            return 0;
        }
        return findNodeCount(info.left.apply(node), info) + findNodeCount(info.right.apply(node), info) + 1;
    }

    @Data
    public static class PrintNodeInfo<T> {
        private T node;
        private Function<T, Integer> val;
        private Function<T, T> left;
        private Function<T, T> right;
        private Function<T, String> valStr;
        private Function<T, Boolean> rb;

        public PrintNodeInfo() {
        }

        public PrintNodeInfo(T node, Function<T, Integer> val, Function<T, T> left, Function<T, T> right) {
            this.node = node;
            this.val = val;
            this.left = left;
            this.right = right;
        }
    }


}

随机生成100000各不同的数字，都符合红黑树。