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各不同的数字，都符合红黑树。
