代码
public class RedBlackTree<Data extends Comparable> {
private static final boolean RED = true;
private static final boolean BLACK = false;
// 空结点
private Node nil = new Node(BLACK);
// 根结点
private Node root = nil;
// 结点数量
private int size;
public int size() {
return size;
}
public class Node {
// 结点数据
private Data data;
// 结点颜色
private boolean color;
// 父结点
private Node parent;
// 左子结点
private Node left;
// 右子结点
private Node right;
public Node(boolean color) {
this.color = color;
}
public Node(Data data, Node parent, Node left, Node right, boolean color) {
this.data = data;
this.parent = parent;
this.left = left;
this.right = right;
this.color = color;
}
@Override
public String toString() {
return "Node{" +
"data=" + data +
", parent=" + (parent != nil ? parent.data : null) +
", left=" + (left != nil ? left.data : null) +
", right=" + (right != nil ? right.data : null) +
", color=" + (color ? "红" : "黑") +
'}';
}
}
/**
* 向树中插入数据
*
* @param data
* @return
*/
public Node insert(Data data) {
if (data == null) {
return null;
}
if (root == nil) {
// 根结点颜色为黑
root = new Node(data, nil, nil, nil, BLACK);
size++;
return root;
}
boolean onLeft = false;
Node currNode = root, tempNode = nil;
while (currNode != nil) {
tempNode = currNode;
int comp = data.compareTo(currNode.data);
if (comp == 0) {
currNode.data = data;
return currNode;
}
onLeft = comp < 0;
currNode = onLeft ? currNode.left : currNode.right;
}
// 创建新结点,颜色默认为红
Node newNode = new Node(data, tempNode, nil, nil, RED);
if (onLeft) {
tempNode.left = newNode;
} else {
tempNode.right = newNode;
}
// 调整树结构
adjust(newNode);
size++;
return newNode;
}
/**
* 调整树结构
*
* @param node
*/
private void adjust(Node node) {
while (node.parent != nil && node.parent.color == RED) {
// 1:若当前是根节点,置颜色为黑
if (node.parent == nil) {
node.color = BLACK;
return;
}
// 2:若父节点是黑,不影响平衡,不需要调整
if (node.parent.color == BLACK) {
return;
}
// 3:若父节点是红 ,违反父子节点不能都为红的原则,需要调整
Node parent = node.parent;
Node grandParent = node.parent.parent;
Node uncle = grandParent.left != parent ? grandParent.left : grandParent.right;
if (uncle != nil && uncle.color == RED) {
// 3.1:若叔结点为红,只需要改变颜色,将叔节点和父节点置黑,祖父节点置红
parent.color = BLACK;
grandParent.color = RED;
uncle.color = BLACK;
// 上一步修改颜色后可能会带来新问题,需要以祖父结点为基础向上依次处理
node = grandParent;
} else {
// 3.2:若叔结点为黑或nil,则需要左旋或右旋,调整树结构
if (grandParent.left == parent) {
if (parent.right == node) {
// 3.2.1:LR(父结点为祖父结点的左子结点,当前结点为父结点的右子结点)
// 需要左旋
node = parent;
leftRotate(node);
}
// 3.2.2:LL(父结点为祖父结点的左子结点,当前结点为父结点的右子结点)
// 先改色(父节点置黑,祖父节点置红),然后右旋
node.parent.color = BLACK;
node.parent.parent.color = RED;
rightRotate(grandParent);
} else {
if (parent.left == node) {
// 3.2.3:RL(父结点为祖父结点的右子结点,当前结点为父结点的左子结点)
// 需要右旋
node = parent;
rightRotate(node);
}
// 3.2.4:RR(父结点为祖父结点的右子结点,当前结点为父结点的右子结点)
// 先改色(父节点置黑,祖父节点置红),然后左旋
node.parent.color = BLACK;
node.parent.parent.color = RED;
leftRotate(grandParent);
}
}
}
root.color = BLACK;
}
/**
* 左旋
*
* @param node
*/
private void leftRotate(Node node) {
Node r = node.right;
if (r.left != nil) {
r.left.parent = node;
}
node.right = r.left;
r.parent = node.parent;
if (node.parent == nil) {
root = r;
} else if (node.parent.left == node) {
node.parent.left = r;
} else {
node.parent.right = r;
}
node.parent = r;
r.left = node;
}
/**
* 右旋
*
* @param node
*/
private void rightRotate(Node node) {
Node l = node.left;
if (l.right != nil) {
l.right.parent = node;
}
node.left = l.right;
l.parent = node.parent;
if (node.parent == nil) {
root = l;
} else if (node.parent.left == node) {
node.parent.left = l;
} else {
node.parent.right = l;
}
node.parent = l;
l.right = node;
}
/**
* 先序遍历
*
* @return
*/
public List<Data> preOrderTreeWalk() {
List<Data> list = new ArrayList<>();
preOrderTreeWalk(root, list);
return list;
}
private void preOrderTreeWalk(Node node, List<Data> list) {
if (node == nil) {
return;
}
// 记录数据
list.add(node.data);
// 递归左子树
preOrderTreeWalk(node.left, list);
// 递归右子树
preOrderTreeWalk(node.right, list);
}
/**
* 层序遍历
*
* @return
*/
public void layerOrderTreeWalk() {
Queue<Node> queue = new ConcurrentLinkedQueue<>();
queue.offer(root);
while (!queue.isEmpty()) {
Node node = queue.poll();
System.out.println(node);
if (node.left != nil) {
queue.offer(node.left);
}
if (node.right != nil) {
queue.offer(node.right);
}
}
}
}
测试
public static void main(String[] args) {
int[] arr = {1, 2, 3, 4, 5, 6, 7, 8, 9, 10};
RedBlackTree<Integer> tree = new RedBlackTree<>();
for (int i = 0; i < arr.length; i++) {
tree.insert(arr[i]);
}
System.out.printf("树中元素数量为:%d\n", tree.size());
System.out.println("层序遍历结果如下:");
tree.layerOrderTreeWalk();
}
输出如下:
树中元素数量为:10
层序遍历结果如下:
Node{data=4, parent=null, left=2, right=6, color=黑}
Node{data=2, parent=4, left=1, right=3, color=黑}
Node{data=6, parent=4, left=5, right=8, color=黑}
Node{data=1, parent=2, left=null, right=null, color=黑}
Node{data=3, parent=2, left=null, right=null, color=黑}
Node{data=5, parent=6, left=null, right=null, color=黑}
Node{data=8, parent=6, left=7, right=9, color=红}
Node{data=7, parent=8, left=null, right=null, color=黑}
Node{data=9, parent=8, left=null, right=10, color=黑}
Node{data=10, parent=9, left=null, right=null, color=红}