红黑树的简单实现（java）

最近研究红黑树，简单的实现了一个java的红黑树代码，亲测没有问题，相关实现的说明都在注释中。

实现时遇到的坑：

实现的时候遇到的坑出现在红黑树的删除阶段，网上各种资料都是说删除的时候按照二叉查找树进行删除就好了，结果这块在进行红黑树平衡的时候，待删除的节点是谁，替换的节点是谁很容易搞混，下面代码在写删除这块代码的时候特意标明了待删除节点和替换的节点是谁，便于在看原理的时候可以更好的理解。

/**
* Created by xiaosong on 2017/12/27.
* 红黑树的简单实现
* 红黑树的特性:
* （1）每个节点或者是黑色，或者是红色。
* （2）根节点是黑色。
* （3）每个叶子节点（NIL）是黑色。 [注意：这里叶子节点，是指为空(NIL或NULL)的叶子节点！]
* （4）如果一个节点是红色的，则它的子节点必须是黑色的。
* （5）从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点。
*/

public class RBTree {
public static class TreeNode {
private T data;//数据值
private TreeNode left;
private TreeNode right;
private TreeNode parent;
private Color color;
private int value;//数据的权重

public T getData() {
return data;
}

public void setData(T data) {
this.data = data;
}

public int getValue() {
return value;
}

public void setValue(int value) {
this.value = value;
}
}

enum Color {
BLACK, RED
}

private TreeNode root;

/**
* 插入节点
* 首先将该节点以二叉平衡树的方式插入到相应位置，然后使其满足红黑树的性质
*
* @param node
*/
public void insert(TreeNode node) {
if (root == null) {//树中无元素，则将插入元素当做根结点
root = node;
root.parent = null;
root.left = null;
root.right = null;
root.color = Color.BLACK;//根节点必是黑色
return;
}
//按二叉查找树的方式将该节点插入
TreeNode comNode = root;
while (comNode != null) {
if (comNode.value >= node.value) {
if (comNode.left == null) {
comNode.left = node;
break;
} else {
comNode = comNode.left;
}
} else {
if (comNode.right == null) {
comNode.right = node;
break;
} else {
comNode = comNode.right;
}
}
}
node.parent = comNode;
node.color = Color.RED;//当前新插入且非根节点初始必为红色
node.left = null;
node.right = null;
doAddBalance(node);//插入节点后，将红黑树进行变换，使其继续符合红黑树的性质
}

/**
* 删除节点
*
* @param node
*/
public void delete(TreeNode node) {
delete(node.value);
}

public void delete(int value) {
//找到要删除的节点
TreeNode needDelete = null;
TreeNode currentNode = root;
while (currentNode != null) {
if (currentNode.value == value) {
needDelete = currentNode;
break;
} else if (currentNode.value > value) {
currentNode = currentNode.left;
} else {
currentNode = currentNode.right;
}
}
if (needDelete != null) {
doDelete(needDelete);
}
}

/**
* 添加时进行的红黑树性质匹配操作
*
* @param node
*/
private void doAddBalance(TreeNode node) {
TreeNode currentNode = node;
TreeNode grandP;
TreeNode uncle;
while (currentNode.parent != null && currentNode.parent.color == Color.RED) {//当前节点的父节点为红色或当前节点为根节点
grandP = currentNode.parent.parent;
if (grandP.left == currentNode.parent) {//当前节点的祖父节点的左孩子是当前节点的父亲节点
uncle = currentNode.parent.parent.right;//叔叔节点
if (uncle == null || uncle.color == Color.BLACK) {//叔叔节点是黑色
if (currentNode.parent.right == currentNode) {//当前节点是右孩子
currentNode = currentNode.parent;//以当前节点父亲节点为新的当前节点
leftRotate(currentNode);//对新的当前节点进行左旋
} else if (currentNode.parent.left == currentNode) {//当前节点是左孩子
currentNode.parent.color = Color.BLACK;//将当前节点的父节点设置为黑色
grandP.color = Color.RED;//将当前节点的祖父节点设置为红色
rightRotate(grandP);//对当前节点的祖父节点进行右旋
}
} else {//叔叔节点是红色
//                    (01) 将“父节点”设为黑色。
currentNode.parent.color = Color.BLACK;
//                    (02) 将“叔叔节点”设为黑色。
uncle.color = Color.BLACK;
//                    (03) 如果祖父节点不是根节点，将“祖父节点”设为“红色”。
if (grandP.parent != null) {
grandP.color = Color.RED;
}
//                    (04) 将“祖父节点”设为“当前节点”(红色节点)；
currentNode = grandP;
}
} else if (grandP.right == currentNode.parent) {//当前节点的祖父节点的右孩子是当前节点的父亲节点
uncle = currentNode.parent.parent.left;//叔叔节点
if (uncle == null || uncle.color == Color.BLACK) {//叔叔节点是黑色
if (grandP.right == currentNode.parent) {//当前节点的祖父节点的右孩子是当前节点的父亲节点
if (currentNode.parent.left == currentNode) {//当前节点是左孩子
currentNode = currentNode.parent;//以当前节点父亲节点为新的当前节点
rightRotate(currentNode);//对新的当前节点进行右旋
} else if (currentNode.parent.right == currentNode) {//当前节点是右孩子
currentNode.parent.color = Color.BLACK;//将当前节点的父节点设置为黑色
grandP.color = Color.RED;//将当前节点的祖父节点设置为红色
leftRotate(grandP);//对当前节点的祖父节点进行左旋
}
}
} else {//叔叔节点是红色
//同上边处理方式相同
currentNode.parent.color = Color.BLACK;
uncle.color = Color.BLACK;
if (grandP.parent != null) {
grandP.color = Color.RED;
}
currentNode = grandP;
}
}
}
}

/**
* 删除时进行的红黑树性质匹配操作
*
* @param node
*/
public void doDeleteBalance(TreeNode node) {
TreeNode currentNode = node;
while (currentNode != null) {
if (currentNode.color == Color.RED) {//当前节点为红色，直接涂黑处理
currentNode.color = Color.BLACK;
return;
} else {//当前节点为黑色，需要分情况处理
if (currentNode.parent == null) return;//根节点，不做处理
TreeNode brother;//当前节点的兄弟节点
brother = getBrother(currentNode);//获取当前节点的兄弟节点
if (node.parent.left == node) {//被删节点是其父节点的左孩子
if (brother.color == Color.RED) {//若当前节点的兄弟节点为红色
currentNode.parent.color = Color.RED;//将当前节点的父节点变为红色
brother.color = Color.BLACK;//将当前节点的兄弟节点置黑
leftRotate(currentNode.parent);//将当前节点的父节点左旋
} else if (brother.color == Color.BLACK) {//若当前节点的兄弟节点为黑色
if ((brother.left == null || brother.left.color == Color.BLACK) && (brother.right == null || brother.right.color == Color.BLACK)) {//兄弟节点的左右孩子都是黑色
brother.color = Color.RED;//当前节点的兄弟节点颜色变为红色
currentNode = currentNode.parent;//将父节点设置为当前节点
} else if (brother.right != null && brother.right.color == Color.RED) {//兄弟节点的右孩子是红色
brother.color = currentNode.parent.color;//将兄弟节点颜色置为当前节点父节点颜色
currentNode.parent.color = Color.BLACK;//当前节点父节点置黑
brother.right.color = Color.BLACK;//兄弟节点右孩子置黑
leftRotate(currentNode.parent);//以当前节点父节点为支点左旋
return;
} else if ((brother.left != null && brother.left.color == Color.RED) && (brother.right == null || brother.right.color == Color.BLACK)) {//兄弟节点的左孩子是红色，右孩子是黑色
brother.left.color = Color.BLACK;//将兄弟节点的左孩子变为黑色
brother.color = Color.RED;//兄弟节点变为红色
rightRotate(brother);//将兄弟节点右旋
}
}
} else if (node.parent.right == node) {//被删节点是其父节点的右孩子
if (brother.color == Color.RED) {//若当前节点的兄弟节点为红色
currentNode.parent.color = Color.RED;//将当前节点的父节点变为红色
brother.color = Color.BLACK;//将当前节点的兄弟节点置黑
rightRotate(currentNode.parent);//将当前节点的父节点右旋
} else if (brother.color == Color.BLACK) {//若当前节点的兄弟节点为黑色
if ((brother.left == null || brother.left.color == Color.BLACK) && (brother.right == null || brother.right.color == Color.BLACK)) {//兄弟节点的左右孩子都是黑色
brother.color = Color.RED;//当前节点的兄弟节点颜色变为红色
currentNode = currentNode.parent;//将父节点设置为当前节点
} else if (brother.left != null && brother.left.color == Color.RED) {//兄弟节点的左孩子是红色
brother.color = currentNode.parent.color;//将兄弟节点颜色置为当前节点父节点颜色
currentNode.parent.color = Color.BLACK;//当前节点父节点置黑
brother.left.color = Color.BLACK;//兄弟节点左孩子置黑
rightRotate(currentNode.parent);//以当前节点父节点为支点右旋
return;
} else if ((brother.right != null && brother.right.color == Color.RED) && (brother.left == null || brother.left.color == Color.BLACK)) {//兄弟节点的右孩子是红色，左孩子是黑色
brother.right.color = Color.BLACK;//将兄弟节点的右孩子变为黑色
brother.color = Color.RED;//兄弟节点变为红色
leftRotate(brother);//将兄弟节点右旋
}
}
}
}
}
}

public TreeNode getBrother(TreeNode node) {//获取当前节点的兄弟节点
if (node.parent.left != node) {//重新设置X的兄弟节点
return node.parent.left;
} else {
return node.parent.right;
}
}

private void doDelete(TreeNode node) {
TreeNode currentNode;//要替换的节点
TreeNode needDelete = node;//要删除的节点
//先进行二叉排序树的删除操作，找到要进行替换的节点
if (node.left != null && node.right != null) {//待删除节点左右子树均在
TreeNode succed = node.left;//寻找该节点的后继节点（找左儿子的最右子树）
while (succed.right != null) {
succed = succed.right;
}
node.value = succed.value;//将原删除节点的值替换为其后继节点的值，将后继节点作为待删除节点
node.data = succed.data;
needDelete = succed;
}
currentNode = needDelete.left != null ? needDelete.left : needDelete.right;
if (currentNode != null) {
currentNode.parent = needDelete.parent;//删除待删除节点，并用其子节点代替
if (needDelete.left == currentNode) {
currentNode.parent.left = currentNode;
} else {
currentNode.parent.right = currentNode;
}
needDelete.left = needDelete.right = needDelete.parent = null;
if (needDelete.color == Color.BLACK) {
doDeleteBalance(node);
}
} else if (needDelete.parent == null) {//说明当前是根节点
root = null;
} else {//说明待删除节点是个叶子节点
currentNode = needDelete;//将自己作为替换节点，同时自己也是待删除节点
if (needDelete.color == Color.BLACK) {
doDeleteBalance(currentNode);
}
//平衡之后，将其删除
if (currentNode.parent != null) {
if (currentNode.parent.left == currentNode) {
currentNode.parent.left = null;
} else if (currentNode.parent.right == currentNode) {
currentNode.parent.right = null;
}
currentNode.parent = null;
}
}
}

public boolean isLeaf(TreeNode node) {
return node.left == null && node.right == null;
}

/**
* 左旋操作，传入的参数为要左旋的节点
*/
private void leftRotate(TreeNode des) {
TreeNode x = des;
TreeNode y = des.right;
TreeNode b = y.left;
x.right = b;//将y的左孩子变为x的右孩子
if (b != null) {
b.parent = x;
}
y.parent = x.parent;//将x的父亲设置为y的父亲
if (x.parent == null) {//根节点
root = y;
} else {//非根节点
if (x.parent.left == x) {//x是其父节点的左孩子
x.parent.left = y;
} else if (x.parent.right == x) {//x是其父节点的右孩子
x.parent.right = y;
}
}
x.parent = y;//将y设置成x的父亲
y.left = x;//将x设置为y的左孩子
}

/**
* 右旋操作，传入的参数为要右旋的节点
*/
private void rightRotate(TreeNode des) {
TreeNode y = des;
TreeNode x = des.left;
TreeNode b = x.right;
y.left = b;//将x的右孩子变为y的左孩子
if (b != null) {
b.parent = y;
}
//将y的父亲设置为x的父亲
x.parent = y.parent;
if (y.parent == null) {//根节点
root = x;
} else {//非根节点
if (y.parent.left == y) {//y是其父节点的左孩子
y.parent.left = x;
} else if (y.parent.right == y) {//y是其父节点的右孩子
y.parent.right = x;
}
}
y.parent = x;//将x设置成y的父亲
x.right = y;//将y设置为x的右孩子
}
}

参考资料：

1.https://www.cnblogs.com/skywang12345/p/3245399.html

2.java TreeMap的源码