通过java简单实现红黑树
package com.example.demo.rbt;
import lombok.Data;
public class RBTree<K extends Comparable<K>, V> {
/**
* 红色
*/
public static final boolean RED = true;
/**
* 黑色
*/
public static final boolean BLACK = false;
/**
* 根节点
*/
private RBNode<K, V> root;
/**
* 获取父节点
* @param node 当前节点
* @return 父节点
*/
private RBNode<K, V> getParentNode(RBNode<K, V> node){
if (node != null) {
return node.parent;
}
return null;
}
/**
* 判断节点是否是红色
* @param node 当前节点
* @return true or false
*/
private boolean isRed(RBNode<K, V> node){
return node != null && node.color;
}
/**
* 判断节点是否是黑色
* @param node 当前节点
* @return true or false
*/
private boolean isBlack(RBNode<K, V> node){
return node != null && !node.color;
}
/**
* 设置为红色
* @param node 当前节点
*/
private void setRed(RBNode<K, V> node){
if (node != null) {
node.setColor(RED);
}
}
/**
* 设置为黑色
* @param node 当前节点
*/
private void setBlack(RBNode<K, V> node){
if (node != null) {
node.setColor(BLACK);
}
}
/**
* 右旋:以某个节点作为支点(旋转节点),其左子节点变为旋转节点的父节点,左子节点的右子节点变为旋转节点的左子节点,右子节点保持不变
* F F
* | |
* K (旋转节点) H
* / \ ==> / \
* H M NIL K
* / \ / \
* NIL NIL NIL M
* @param node 旋转节点
*/
private void rightSpin(RBNode<K, V> node){
// 当前节点的父节点 F
RBNode<K, V> pN = node.parent;
// 当前节点的左子节点 H
RBNode<K, V> lN = node.left;
// 将K的左子节点指向H的右子节点NIL,将H的右子节点的父节点更新为K K--NIL
node.left = lN.right;
if (lN.right != null) {
lN.right.parent = node;
}
// 将K的父节点更新为H,将H的右子节点更新为K H--K
node.parent = lN;
lN.right = node;
// 当F不为空时(如果K是根节点就有可能为null),将H的父节点更新为F,将F的左(右)子节点更新为H H--F
if (pN != null) {
lN.parent = pN;
if (pN.left == node) {
pN.left = lN;
}else {
pN.right = lN;
}
}else {
// 将H设置为根节点
this.root = lN;
}
}
/**
* 左旋:以某个节点作为支点(旋转节点),其右子结点变为旋转节点的父节点,右子结点的左子节点变为旋转节点的右子节点,左子节点保持不变
* F F
* | |
* K (旋转节点) M
* / \ ==> / \
* H M K NIL
* / \ / \
* NIL NIL H NIL
* @param node 旋转节点
*/
private void leftSpin(RBNode<K, V> node){
// 当前节点的父节点 F
RBNode<K, V> pN = node.parent;
// 当前节点的右子节点 M
RBNode<K, V> rN = node.right;
// 将K的右子节点指向M的左子节点NIL,将M的左子节点的父节点更新为K K--NIL
node.right = rN.left;
if (rN.left != null) {
rN.left.parent = node;
}
// 将K的父节点更新为M,将M的左子节点更新为K K--M
node.parent = rN;
rN.left = node;
// 当F不为空时(如果K是根节点就有可能为null),将M的父节点更新为F,将F的左(右)子节点更新为M M--F
rN.parent = pN;
if (pN != null) {
// 如果K原先是左节点则将F的左节点更新为M,如果是右节点则将F的右节点更新为M
if (pN.left == node) {
pN.left = rN;
}else {
pN.right = rN;
}
}else {
// 将M设置为根节点
this.root = rN;
}
}
/**
* 中序遍历
*/
public void inOrderPrint(){
System.out.print("中序遍历:");
inOrderPrint(this.root);
}
public RBNode<K, V> getRoot(){
return this.root;
}
private void inOrderPrint(RBNode<K, V> node){
if (node != null) {
inOrderPrint(node.left);
System.out.print(node.key + " ");
inOrderPrint(node.right);
}
}
public RBNode<K, V> get(K key){
return getNode(this.root, key);
}
private RBNode<K, V> getNode(RBNode<K, V> node, K key){
while (node != null){
K k = node.key;
if (key.compareTo(k) > 0) {
node = node.right;
}else if(key.compareTo(k) < 0){
node = node.left;
}else {
return node;
}
}
return null;
}
public void add(K key, V value){
RBNode<K, V> node = new RBNode<>(RED, key, value);
addNode(node);
}
private void addNode(RBNode<K, V> node){
// 查找当前节点的父节点,从root开始
RBNode<K, V> parent = null;
RBNode<K, V> temp = this.root;
while (temp != null){
parent = temp;
K k = temp.key;
int i = node.key.compareTo(k);
if (i > 0) {
// 插入节点的Key大于当前节点的Key,从当前节点的右子树找
temp = temp.right;
} else if (i < 0) {
// 插入节点的Key小于当前节点的Key,从当前节点的左子树找
temp = temp.left;
} else {
// 插入节点的Key等于当前节点的Key,直接替换value TODO 场景二
temp.value = node.value;
return;
}
}
// 设置插入节点的父节点
node.parent = parent;
if (parent != null) {
int i = node.key.compareTo(parent.key);
if (i > 0) {
// 如果插入节点的Key大于父节点的Key,则放在父节点的右子节点上
parent.right = node;
}else {
// 如果插入节点的Key小于父节点的Key,则放在父节点的左子节点上
parent.left = node;
}
}else {
// 说明是颗空树,将插入节点设置为根节点 TODO 场景一
this.root = node;
}
// 插入完后走自平衡处理
selfBalance(node);
}
/**
* 插入节点的后续处理
* |--- 1.红黑树为空数,将插入节点设置为根节点并置为黑色
* |--- 2.插入节点的Key已存在,则更新已存在节点的Value为插入节点的Value
* |--- 3.插入节点的父节点是黑色的,则直接插入,不会影响平衡
* |--- 4.插入节点的父节点是红色的
* |--- 4.1 叔叔节点存在并且为红节点
* - 将父节点(F)和叔叔节点(V)设置为黑色
* - 将祖父节点(P)设置为红色
* - 将祖父节点设置为当前插入节点
* |--- 4.2 叔叔节点不存在或者为黑节点,并且插入节点的父节点是祖父节点的左子节点
* |--- 4.2.1 插入节点是其父节点的左子节点(LL)
* - 将父节点(F)设置为黑色
* - 将祖父节点(P)设置为红色
* - 对祖父节点(P)进行右旋
* |--- 4.2.1 插入节点是其父节点的右子节点(LR)
* - 对父节点(F)进行左旋
* - 把父节点(F)设置为插入结点,得到情景4.2.1
* - 进行情景4.2.1的处理
* |--- 4.3 叔叔节点不存在或者为黑节点,并且插入节点的父节点是祖父节点的右子节点
* |--- 4.3.1 插入节点是其父节点的右子节点(RR)
* - 将父节点(F)设置为黑色
* - 将祖父节点(P)设置为红色
* - 对祖父节点(P)进行左旋
* |--- 4.3.2 插入节点是其父节点的左子节点(RL)
* - 对父节点(F)进行右旋
* - 把父节点(F)设置为插入结点,得到情景4.3.1
* - 进行情景4.3.1的处理
* @param node
*/
private void selfBalance(RBNode<K, V> node){
// 1.红黑树为空数,将插入节点设置为根节点并置为黑色
if (this.root == node) {
setBlack(node);
return;
}
// 2.插入节点的Key已存在,则更新已存在节点的Value为插入节点的Value
// 3.插入节点的父节点是黑色的,则直接插入,不会影响平衡
// 4.插入节点的父节点是红色的
// 父节点
RBNode<K, V> pN = getParentNode(node);
// 祖父节点
RBNode<K, V> gpN = getParentNode(pN);
if (isRed(pN)) {
// 叔叔节点
RBNode<K, V> uN;
if (gpN.left == pN) {
uN = gpN.right;
}else {
uN = gpN.left;
}
// 4.1 叔叔节点存在并且为红节点
if (isRed(uN)) {
// 将父节点和叔叔节点设置为黑色,祖父节点设置为红色,然后将祖父节点设置为当前节点进行递归
setBlack(pN);
setBlack(uN);
setRed(gpN);
selfBalance(gpN);
return;
}
// 4.2 叔叔节点不存在或者为黑节点
if (uN == null || isBlack(uN)) {
// 插入节点的父节点是祖父节点的左子节点
if (pN == gpN.left) {
// 4.2.1 插入节点是其父节点的左子节点(LL)
if (node == pN.left) {
setBlack(pN);
setRed(gpN);
rightSpin(gpN);
}else {
// 4.2.2 插入节点是其父节点的右子节点(LR)
leftSpin(pN );
selfBalance(pN);
}
}else {
// 插入节点的父节点是祖父节点的右子节点
if (node == node.parent.left) {
// 4.3.2 插入节点是其父节点的左子节点(RL)
rightSpin(pN);
selfBalance(pN);
}else {
// 4.3.1 插入节点是其父节点的右子节点(RR)
setBlack(pN);
setRed(gpN);
leftSpin(gpN);
}
}
}
}
}
@Data
static class RBNode<K extends Comparable<K>, V>{
/**
* 父节点
*/
private RBNode<K, V> parent;
/**
* 右子节点
*/
private RBNode<K, V> right;
/**
* 左子节点
*/
private RBNode<K, V> left;
/**
* 颜色
*/
private boolean color;
/**
* 键
*/
private K key;
/**
* 值
*/
private V value;
public RBNode() {
}
public RBNode(boolean color, K key, V value) {
this.color = color;
this.key = key;
this.value = value;
}
/**
* 这里重写toString方法是因为子节点包含父节点的引用,父节点也包含子节点的引用,
* 在debug的时候会调用toString方法显示结构,相互引用造成StackOverflowError
*/
@Override
public String toString() {
return "RBNode{" +
"color=" + color +
", key=" + key +
", value=" + value +
'}';
}
}
}
打印红黑树结构(网上找的一个方法)
package com.example.demo.rbt;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Random;
public class RBTreePrinterTest {
public static <K extends Comparable<K>, V> void printNode(RBTree.RBNode<K, V> root) {
int maxLevel = RBTreePrinterTest.maxLevel(root);
printNodeInternal(Collections.singletonList(root), 1, maxLevel);
}
private static <K extends Comparable<K>, V> void printNodeInternal(List<RBTree.RBNode<K, V>> nodes, int level, int maxLevel) {
if (nodes.isEmpty() || RBTreePrinterTest.isAllElementsNull(nodes)) {
return;
}
int floor = maxLevel - level;
int endgeLines = (int) Math.pow(2, (Math.max(floor - 1, 0)));
int firstSpaces = (int) Math.pow(2, (floor)) - 1;
int betweenSpaces = (int) Math.pow(2, (floor + 1)) - 1;
RBTreePrinterTest.printWhitespaces(firstSpaces);
List<RBTree.RBNode<K, V>> newNodes = new ArrayList<>();
for (RBTree.RBNode<K, V> node : nodes) {
if (node != null) {
String key = node.isColor() ? "\033[31;4m" + node.getKey() + "\033[0m" : node.getKey().toString();
System.out.print(key);
newNodes.add(node.getLeft());
newNodes.add(node.getRight());
} else {
newNodes.add(null);
newNodes.add(null);
System.out.print(" ");
}
RBTreePrinterTest.printWhitespaces(betweenSpaces);
}
System.out.println();
for (int i = 1; i <= endgeLines; i++) {
for (RBTree.RBNode<?, ?> node : nodes) {
RBTreePrinterTest.printWhitespaces(firstSpaces - i);
if (node == null) {
RBTreePrinterTest.printWhitespaces(endgeLines + endgeLines + i + 1);
continue;
}
if (node.getLeft() != null) {
System.out.print("/");
} else {
RBTreePrinterTest.printWhitespaces(1);
}
RBTreePrinterTest.printWhitespaces(i + i - 1);
if (node.getRight() != null) {
System.out.print("\\");
} else {
RBTreePrinterTest.printWhitespaces(1);
}
RBTreePrinterTest.printWhitespaces(endgeLines + endgeLines - i);
}
System.out.println("");
}
printNodeInternal(newNodes, level + 1, maxLevel);
}
private static void printWhitespaces(int count) {
for (int i = 0; i < count; i++) {
System.out.print(" ");
}
}
private static <K extends Comparable<K>, V> int maxLevel(RBTree.RBNode<K, V> node) {
if (node == null) {
return 0;
}
return Math.max(RBTreePrinterTest.maxLevel(node.getLeft()), RBTreePrinterTest.maxLevel(node.getRight())) + 1;
}
private static <K extends Comparable<K>, V> boolean isAllElementsNull(List<RBTree.RBNode<K, V>> list) {
for (Object object : list) {
if (object != null) {
return false;
}
}
return true;
}
public static void main(String[] args) {
RBTree<Integer, Object> rbTree = new RBTree<>();
// 随机15位数插入
new Random().ints(15, 1, 50).forEach(k -> rbTree.add(k, null));
RBTreePrinterTest.printNode(rbTree.getRoot());
rbTree.inOrderPrint();
}
}
实现效果
红色字体需要安装IDEA插件 ANSI Highlighter