红黑树的插入-java代码实现

本文基于上一博客红黑树的插入理论讲解-多图进行代码实现，若代码有疑问，可以结合上文。

插入调整Case汇总

第一种调整：CASEA1 + CASEB1，算法一致。新增节点可出现在下列四个任意位置

第二种调整：CASEA2 + CASEA3，若出现CASEA2 ，先调整为CASEA3，再最终调整为红黑树子树

第三种调整：CASEB2 + CASEB3，若出现CASEB2 ，先调整为CASEB3，再最终调整为红黑树子树

实现

针对CASE 进行代码实现

定义颜色

public enum Color {
	BLACK, RED;
}

定义节点

public class RBNode {
	private int value;
	private RBNode left;
	private RBNode right;
	private RBNode parent;
	private Color color = Color.BLACK;
	// setter getter
}

平衡二叉树插入

public class RBTree {
	
	private RBNode root;
	
	public void put(int value) {
		if (root == null) {
			root = new RBNode();
			root.setValue(value);
			return;
		}
		
		RBNode point = root;
		RBNode parent = null;
		do {
			parent = point;
			int pointValue = point.getValue();
			if (pointValue == value) {
				return;
			}
			if (value < pointValue) {
				point = point.getLeft();
			} else {
				point = point.getRight();
			}
		} while (point != null);
		
		// 新插入的结点都是红色
		RBNode node = new RBNode();
		node.setValue(value);
		node.setColor(Color.RED);
		
		if (value < parent.getValue()) {
			parent.setLeft(node);
		} else {
			parent.setRight(node);
		}
		node.setParent(parent);
		// 涉及递归，判断只能在方法内做
		reBalanceAfterInsert(node);
	}
	
	private RBNode parentOf(RBNode p) {
		return (p == null ? null : p.getParent());
	}

	private RBNode leftOf(RBNode p) {
		return (p == null) ? null : p.getLeft();
	}

	private RBNode rightOf(RBNode p) {
		return (p == null) ? null : p.getRight();
	}

	
	private void reBalanceAfterInsert(RBNode node) {
		if (parentOf(node) != null && parentOf(node).getColor() == Color.RED) {
			// 一定存在，且是红色
			RBNode parent = parentOf(node);
			// 一定存在，且是黑色
			RBNode grandpa = parentOf(parent);
			
			RBNode leftParent = leftOf(grandpa);
			RBNode rightParent = rightOf(grandpa);
			
			// CASEA1 + CASEB1
			if (leftParent != null && rightParent != null && leftParent.getColor() == Color.RED
					&& rightParent.getColor() == Color.RED) {
				grandpa.setColor(Color.RED);
				leftParent.setColor(Color.BLACK);
				rightParent.setColor(Color.BLACK);
				// 递归，黑变红可能出现红红，使用祖父结点重新判断
				if (node != root) {
					reBalanceAfterInsert(grandpa);
				}
			} 
			// CASEA2 + CASEA3
			else if (parent == leftParent) {
				// CASEA2
				if (rightOf(parent) == node) {
					node = parent;
					rotateLeft(node);
				}
				// CASEA3
				parentOf(node).setColor(Color.BLACK);
				parentOf(parentOf(node)).setColor(Color.RED);
				rotateRight(parentOf(parentOf(node)));
				// end
			} 
			// CASEB2 + CASEB3
			else if (parent == rightParent) {
				// CASEB2
				if (leftOf(parent) == node) {
					node = parent;
					rotateRight(node);
				}
				// CASEB3
				parentOf(node).setColor(Color.BLACK);
				parentOf(parentOf(node)).setColor(Color.RED);
				rotateLeft(parentOf(parentOf(node)));
				// end
			}
		}
		root.setColor(Color.BLACK);

	}

	// 左旋-镜像对称右旋
	private void rotateLeft(RBNode x) {
		RBNode y = x.getRight();
		x.setRight(y.getLeft());
		if (y.getLeft() != null) {
			y.getLeft().setParent(x);
		}
		
		RBNode parent = x.getParent();
		if (parent == null) {
			// 旋转点可能根节点，需要切换
			root = y;
		} else if (leftOf(parent) == x) {
			parent.setLeft(y);
		} else {
			parent.setRight(y);
		}
		y.setParent(parent);
		
		y.setLeft(x);
		x.setParent(y);
	}
	
	// 右旋-镜像对称左旋
	private void rotateRight(RBNode y) {
		RBNode x= y.getLeft();
		y.setLeft(x.getRight());
		if (x.getRight() != null) {
			x.getRight().setParent(y);
		}
		
		RBNode parent = y.getParent();
		if (parent == null) {
			// 旋转点可能根节点，需要切换
			root = x;
		} else if (leftOf(parent) == y) {
			parent.setLeft(x);
		} else {
			parent.setRight(x);
		}
		x.setParent(parent);
		
		x.setRight(y);
		y.setParent(x);
	}
	
	// 层序遍历输出
	public void print() {
		Queue<RBNode> queue = new LinkedBlockingDeque<>();
		queue.add(root);

		while (!queue.isEmpty()) {
			RBNode pollNode = queue.poll();
			System.out.println(String.format("value=%s, color=%s", pollNode.getValue(), pollNode.getColor()));
			RBNode leftNode = pollNode.getLeft();
			if (leftNode != null) {
				queue.add(leftNode);
			}
			RBNode rightNode = pollNode.getRight();
			if (rightNode != null) {
				queue.add(rightNode);
			}
		}
	}
}

测试

public class RBTreeTest {
	public static void main(String[] args) {
		int nums[] = { 4, 1, 2, 5, 0, 7, 9, 8, 3, 6};
		RBTree tree = new RBTree();
		for (int i : nums) {
			tree.put(i);
		}
		
		tree.print();
	}
}

输出

value=5, color=BLACK
value=2, color=RED
value=8, color=RED
value=1, color=BLACK
value=4, color=BLACK
value=7, color=BLACK
value=9, color=BLACK
value=0, color=RED
value=3, color=RED
value=6, color=RED