红黑树是一种高效平衡二叉树,实现起来稍微复杂。
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.Set;
//假设键不重复
public class RBTree<T extends Comparable<T>> {
public Node<T> root;
static final boolean Red = false;
static final boolean Black = true;
class Node<T extends Comparable<T>> {
Node<T> parent;
Node<T> left;
Node<T> right;
T data;
boolean color;
Node(T data, boolean color, Node<T> parent, Node<T> left, Node<T> right) {
this.data = data;
this.color = color;
this.parent = parent;
this.left = left;
this.right = right;
}
}
// 辅助函数:左旋、右旋某个节点(此节点是待旋转子树的root)
void leftRotate(Node<T> node) {
Node<T> temp = node.right;
node.right = temp.left;
if (node.right != null) {
node.right.parent = node;
}
temp.parent = node.parent;
if (temp.parent != null) {
if (node.parent.left == node)
temp.parent.left = temp;
else
temp.parent.right = temp;
} else
root = temp;
temp.left = node;
node.parent = temp;
}
void rightRotate(Node<T> node) {
Node<T> temp = node.left;
node.left = temp.right;
if (node.left != null) {
node.left.parent = node;
}
temp.parent = node.parent;
if (temp.parent != null) {
if (node.parent.left == node)
temp.parent.left = temp;
else
temp.parent.right = temp;
} else
root = temp;
temp.right = node;
node.parent = temp;
}
public boolean find(T data) {
Node<T> current = root;
while (current != null) {
if (current.data.compareTo(data) == 0)
return true;
else if (data.compareTo(current.data) > 0)
current = current.right;
else
current = current.left;
}
return false;
}
// 先找到插入位置,再调整
private void insert(Node<T> node) {
Node<T> current = root;
Node<T> temp = null;
int cmp;
while (current != null) {
temp = current;// 保存插入位置的父节点
cmp = node.data.compareTo(current.data);
if (cmp < 0)
current = current.left;
else
current = current.right;
}
if (temp == null)// while第一次就失败,说明root为null
root = node;
else {
cmp = node.data.compareTo(temp.data);
if (cmp < 0) {
temp.left = node;
} else {
temp.right = node;
}
}
node.parent = temp;
root.color = Black;// 根为黑
// 插入完后调整结构
fixAfterInsert(node);
}
// 包装值后再插入
public void insert(T data) {
Node<T> node = new Node<T>(data, Red, null, null, null);// 新插入的都设置为红,调整起来便利
if (node != null)
insert(node);
}
private void fixAfterInsert(Node<T> node) {
// 小于3层不用操作
while (node.parent != null && node.parent.parent != null && node.parent.color == Red) {
// 左外侧
if (node.parent == node.parent.parent.left && node == node.parent.left) {
// 如果没有叔节点(只可能有一个红色叔节点)旋转一次,无需向上调整
if (node.parent.parent.right == null) {
node.parent.color = Black;
node.parent.parent.color = Red;
rightRotate(node.parent.parent);
break;
} else {
node.color = Black;
rightRotate(node.parent.parent);
node = node.parent;
continue;
}
}
// 右外侧
else if (node.parent == node.parent.parent.right && node == node.parent.right) {
if (node.parent.parent.left == null) {
node.parent.color = Black;
node.parent.parent.color = Red;
leftRotate(node.parent.parent);
break;
} else {
node.color = Black;
leftRotate(node.parent.parent);
node = node.parent;
continue;
}
}
// 左-右
else if (node.parent == node.parent.parent.left && node == node.parent.right) {
if (node.parent.parent.right == null) {
node.color = Black;
node.parent.parent.color = Red;
leftRotate(node.parent);
rightRotate(node.parent);
break;
} else {
node.parent.color = Black;
leftRotate(node.parent);
rightRotate(node.parent);
continue;
}
}
// 右-左
else if (node.parent == node.parent.parent.right && node == node.parent.left) {
if (node.parent.parent.left == null) {
node.color = Black;
node.parent.parent.color = Red;
rightRotate(node.parent);
leftRotate(node.parent);
break;
} else {
node.parent.color = Black;
rightRotate(node.parent);
leftRotate(node.parent);
continue;
}
}
}
}
// 从小到大遍历来一个
public void display(Node<T> root) {
if (root == null)
return;
display(root.left);
System.out.print(root.data + "/");
display(root.right);
}
// 用于显示数据所处位置和颜色,data为数值,i为data插入的行,list为方向变化,1为向右,-1向左
// b为data的颜色,true为黑
void test(T data, Node<T> rt) {
int i = 0;// i记录行
boolean b;
List<Integer> list = new ArrayList<>();
Node<T> current = rt;
while (current != null) {
if (data.compareTo(current.data) < 0) {
current = current.left;
i++;// 行
list.add(-1);
} else if (data.compareTo(current.data) > 0) {
current = current.right;
i++;
list.add(1);
} else if (data.compareTo(current.data) == 0) {
b = current.color;
System.out.println(i + "行-->" + data + "-->" + list + "-->" + b);
break;
}
}
}
// 测试是否有2个红色节点接连出现
void test1(Node<T> node) {
if (node == null)
return;
if (node.color == Red
&& ((node.left != null && node.left.color == Red) || (node.right != null && node.right.color == Red))) {
System.out.println("two red occur");
return;
}
test1(node.right);
test1(node.left);
}
public static void main(String[] args) {
RBTree<Integer> tree = new RBTree<>();
// 用set测试不同值
Set<Integer> set = new HashSet<>();
Random rand = new Random();
for (int i = 0; i < 50; i++)
set.add(rand.nextInt(100));
List<Integer> list = new ArrayList<>(set);
Iterator<Integer> it = set.iterator();
while (it.hasNext())
tree.insert(it.next());
// tree.display(tree.root);
tree.test1(tree.root);
for (Integer i : list)
tree.test(i, tree.root);
}
}