Skip to content

Instantly share code, notes, and snippets.

@geminiwen
Last active September 30, 2022 06:21
Show Gist options
  • Save geminiwen/82302e0d743ac5289fedd3821b4d8bcc to your computer and use it in GitHub Desktop.
Save geminiwen/82302e0d743ac5289fedd3821b4d8bcc to your computer and use it in GitHub Desktop.
RBTree
package com.gemini.demo;
public class Main {
public static void main(String[] args) {
// write your code here
TreeNode root = new TreeNode();
root.setColor(TreeNode.COLOR_BLACK);
root.setData(50);
root = root.insert(10);
root = root.insert(20);
root = root.insert(30);
root = root.insert(40);
root = root.insert(50);
root = root.insert(60);
root = root.insert(70);
root = root.insert(80);
root = root.insert(90);
root = root.insert(100);
System.out.println(root);
}
}
package com.gemini.demo;
public class TreeNode {
public static final int COLOR_RED = 0;
public static final int COLOR_BLACK = 1;
private TreeNode left;
private TreeNode right;
private TreeNode parent;
private int color;
private int data;
public TreeNode() {
this.color = COLOR_RED;
this.left = null;
this.right = null;
this.parent = null;
}
public TreeNode getLeft() {
return left;
}
public void setLeft(TreeNode left) {
this.left = left;
}
public TreeNode getRight() {
return right;
}
public void setRight(TreeNode right) {
this.right = right;
}
public int getColor() {
return color;
}
public void setColor(int color) {
this.color = color;
}
public int getData() {
return data;
}
public void setData(int data) {
this.data = data;
}
public TreeNode grandParent() {
return this.parent.parent;
}
public TreeNode uncle() {
if (this.parent == this.grandParent().left) {
return this.grandParent().right;
} else {
return this.grandParent().left;
}
}
/**
* 插入一个值,返回 Root
* @return
*/
public TreeNode insert(int n) {
TreeNode child = TreeNode.createNode(this, n);
if (child != null) {
insertNode(child);
}
return this.findRoot();
}
/**
* 把 node 加入到这棵树中去
* @param child 子节点
* @return 根节点
*/
public void insertNode(TreeNode child) {
TreeNode parent = child.parent;
if (parent == null) {
child.color = COLOR_BLACK;
return;
}
// 如果父亲的节点为黑色
if (parent.color == COLOR_BLACK) {
return;
}
// 如果父亲节点为红色,然后叔叔节点的颜色也为红色
if (child.uncle() != null &&
child.uncle().color == COLOR_RED) {
child.parent.color = COLOR_BLACK;
child.uncle().color = COLOR_BLACK;
child.grandParent().color = COLOR_RED;
insertNode(child.grandParent());
} else {
rotateInsertNode(child);
}
}
private void rotateInsertNode(TreeNode child) {
TreeNode node = child;
if (child == child.parent.right &&
child.parent == child.grandParent().left) {
node = child.parent.rotateLeft();
node = node.left;
} else if (child == child.parent.left &&
child.parent == child.grandParent().right) {
node = child.parent.rotateRight();
node = node.right;
}
checkResult(node);
}
private TreeNode rotateLeft() {
TreeNode nextParent = this.right;
this.right = this.right.left;
nextParent.left = this;
if (this.parent != null) {
if (this.parent.left == this) {
this.parent.left = nextParent;
} else {
this.parent.right = nextParent;
}
}
nextParent.parent = this.parent;
this.parent = nextParent;
return nextParent;
}
private TreeNode rotateRight() {
TreeNode nextParent = this.left;
this.left = nextParent.right;
nextParent.right = this;
if (this.parent != null) {
if (this.parent.right == this) {
this.parent.right = nextParent;
} else {
this.parent.left = nextParent;
}
}
nextParent.parent = this.parent;
this.parent = nextParent;
return nextParent;
}
/**
* 检查阶段,看是否需要旋转
* @param child
* @return
*/
private void checkResult(TreeNode child) {
child.parent.color = COLOR_BLACK;
child.grandParent().color = COLOR_RED;
if (child == child.parent.left && child.parent == child.grandParent().left) {
child.grandParent().rotateRight();
} else {
child.grandParent().rotateLeft();
}
}
private static TreeNode createNode(TreeNode root, int n) {
TreeNode parent = root.findParentForNewValue(n);
// 待插入的节点已经存在
if (parent == null) return null;
// 我们设定插入的孩子必须永远为红色
TreeNode child = new TreeNode();
child.color = COLOR_RED;
child.data = n;
child.parent = parent;
if (parent.data > n) {
parent.left = child;
} else {
parent.right = child;
}
return child;
}
/**
* 为新插入的值找一个合适的 parent
* @param n 值
* @return 新插入的n值的父亲
*/
private TreeNode findParentForNewValue(int n) {
if (this.data == n) {
return null;
} else if (this.data > n) {
if (this.left == null) return this;
return this.left.findParentForNewValue(n);
} else {
if (this.right == null) return this;
return this.right.findParentForNewValue(n);
}
}
private TreeNode findRoot() {
TreeNode node = this;
while (node.parent != null) {
node = node.parent;
}
return node;
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment