public class RedBlackTree<K extends Comparable<K>, V> {
private static final boolean RED = false;
private static final boolean BLACK = true;
private TreeNode<K, V> root;
public TreeNode<K, V> getRoot() {
return root;
}
public void setRoot(TreeNode<K, V> root) {
this.root = root;
}
static class TreeNode<K, V> {
private TreeNode<K, V> parent;
private TreeNode<K, V> left;
private TreeNode<K, V> right;
private boolean color = BLACK;
private K key;
private V value;
public TreeNode(TreeNode<K, V> parent, K key, V value) {
this.parent = parent;
this.key = key;
this.value = value;
}
public TreeNode(TreeNode<K, V> parent, TreeNode<K, V> left, TreeNode<K, V> right, boolean color, K key, V value) {
this.parent = parent;
this.left = left;
this.right = right;
this.color = color;
this.key = key;
this.value = value;
}
}
/**
* 新增结点
* 1. 普通二叉树的插入
* ① 先查找插入的位置
* ② 将 K,V 封装成TreeNode对象,插入到红黑树中
* 2. 红黑树的平衡(旋转+变色)
*/
public void put(K key, V value) {
if (key == null) {
throw new NullPointerException();
}
TreeNode<K, V> t = this.root;
if (t == null) {
// 红黑树为空,直接创建结点作为根节点
this.root = new TreeNode<>(null, key, value);
setColor(root, BLACK);
return;
}
// 查找插入的位置
int cmp;
// 记录寻找结点的父节点
TreeNode<K, V> parent;
do {
parent = t;
cmp = key.compareTo(t.key);
if (cmp < 0) {
t = t.left;
} else if (cmp > 0) {
t = t.right;
} else {
t.value = value;
return;
}
} while (t != null);
TreeNode<K, V> node = new TreeNode<>(parent, key, value);
if (cmp < 0) {
parent.left = node;
} else {
parent.right = node;
}
// 平衡
balanceInsertion(node);
}
public V remove(K key) {
TreeNode<K, V> removed = find(key);
if (removed == null) {
return null;
}
V oldValue = removed.value;
delete(removed);
return oldValue;
}
private void delete(TreeNode<K, V> removed) {
if (leftOf(removed) != null && rightOf(removed) != null) {
TreeNode<K, V> successor = successor(removed);
removed.key = successor.key;
removed.value = successor.value;
removed = successor;
}
TreeNode<K, V> replacement = removed.left != null ? removed.left : removed.right;
if (replacement != null) {
replacement.parent = removed.parent;
if (removed.parent == null) {
root = replacement;
} else if (removed == leftOf(parentOf(removed))) {
parentOf(removed).left = replacement;
} else {
parentOf(removed).right = replacement;
}
removed.left = removed.right = removed.parent = null;
if (removed.color == BLACK) {
balanceDeletion(replacement);
}
} else if (removed.parent == null) {
root = null;
} else {
if (removed.color == BLACK) {
balanceDeletion(removed);
}
if (removed == leftOf(parentOf(removed))) {
parentOf(removed).left = null;
} else {
parentOf(removed).right = null;
}
removed = null;
}
}
private TreeNode<K, V> find(K key) {
TreeNode<K, V> r = this.root;
while (r != null) {
int cmp = key.compareTo(r.key);
if (cmp < 0) {
r = r.left;
} else if (cmp > 0) {
r = r.right;
} else {
return r;
}
}
return null;
}
private TreeNode<K, V> predecessor(TreeNode<K, V> node) {
if (node == null) {
return null;
} else if (node.left != null) {
TreeNode<K, V> p = node.left;
while (p.right != null) {
p = p.right;
}
return p;
} else {
TreeNode<K, V> parent = node.parent;
TreeNode<K, V> child = node;
while (parent != null && child == parent.left) {
child = parent;
parent = parent.parent;
}
return parent;
}
}
private TreeNode<K, V> successor(TreeNode<K, V> node) {
if (node == null) {
return null;
} else if (node.right != null) {
TreeNode<K, V> p = node.left;
while (p.left != null) {
p = p.left;
}
return p;
} else {
TreeNode<K, V> parent = node.parent;
TreeNode<K, V> child = node;
while (parent != null && child == parent.right) {
child = parent;
parent = parent.parent;
}
return parent;
}
}
/**
* 左旋的实现
* p r
* / \ / \
* l r ===> p rr
* / \ / \
* rl rr l rl
*/
private void rotateLeft(TreeNode<K, V> p) {
TreeNode<K, V> r, pp, rl;
if (p != null && (r = p.right) != null) {
if ((rl = p.right = r.left) != null) {
rl.parent = p;
}
if ((pp = r.parent = p.parent) == null) {
(root = r).color = BLACK;
} else if (pp.left == p) {
pp.left = r;
} else {
pp.right = r;
}
r.left = p;
p.parent = r;
}
}
/**
* 右旋的实现
* p l
* / \ / \
* l r ===> ll p
* / \ / \
* ll rr lr r
*/
private void rotateRight(TreeNode<K, V> p) {
TreeNode<K, V> l, pp, lr;
if (p != null && (l = p.left) != null) {
if ((lr = p.left = l.right) != null) {
lr.parent = p;
}
if ((pp = l.parent = p.parent) == null) {
(root = l).color = BLACK;
} else if (pp.right == p) {
pp.right = l;
} else {
pp.left = l;
}
l.right = p;
p.parent = l;
}
}
/**
* 处理颜色,叶结点是黑色
*/
private boolean colorOf(TreeNode<K, V> node) {
return node == null ? BLACK : node.color;
}
private TreeNode<K, V> parentOf(TreeNode<K, V> node) {
return node == null ? null : node.parent;
}
private TreeNode<K, V> leftOf(TreeNode<K, V> node) {
return node == null ? null : node.left;
}
private TreeNode<K, V> rightOf(TreeNode<K, V> node) {
return node == null ? null : node.right;
}
private void setColor(TreeNode<K, V> node, boolean color) {
if (node != null) {
node.color = color;
}
}
/**
* 插入结点后平衡
*/
private void balanceInsertion(TreeNode<K, V> insert) {
// 插入的结点都是红色结点
setColor(insert, RED);
TreeNode<K, V> parent, uncle, grandparent;
while (insert != null && insert != root && parentOf(insert).color == RED) {
parent = parentOf(insert);
grandparent = parentOf(parent);
// insert 的父结点是爷爷结点的左节点
if (parent == leftOf(grandparent)) {
uncle = rightOf(grandparent);
if (colorOf(uncle) == RED) {
setColor(parent, BLACK);
setColor(uncle, BLACK);
setColor(grandparent, RED);
insert = grandparent;
} else {
if (insert == rightOf(parent)) {
insert = parent;
rotateLeft(insert);
}
setColor(parent, BLACK);
setColor(grandparent, RED);
rotateRight(grandparent);
}
}
// insert 的父结点是爷爷结点的右节点
else {
uncle = leftOf(grandparent);
if (colorOf(uncle) == RED) {
setColor(parent, BLACK);
setColor(uncle, BLACK);
setColor(grandparent, RED);
insert = grandparent;
} else {
if (insert == leftOf(parent)) {
insert = parent;
rotateRight(insert);
}
setColor(parent, BLACK);
setColor(grandparent, RED);
rotateLeft(grandparent);
}
}
}
// 根节点都为黑色结点
setColor(root, BLACK);
}
private void balanceDeletion(TreeNode<K, V> removed) {
if (removed != root && colorOf(removed) == BLACK) {
if (removed == leftOf(parentOf(removed))) {
TreeNode<K, V> sibling = rightOf(parentOf(removed));
if (colorOf(sibling) == RED) {
setColor(sibling, BLACK);
setColor(parentOf(removed), RED);
rotateLeft(parentOf(removed));
sibling = rightOf(parentOf(removed));
}
if (colorOf(leftOf(sibling)) == BLACK && colorOf(parentOf(sibling)) == BLACK) {
setColor(sibling, RED);
removed = parentOf(removed);
} else {
if (colorOf(rightOf(sibling)) == BLACK) {
setColor(sibling, RED);
setColor(leftOf(sibling), BLACK);
rotateRight(sibling);
sibling = rightOf(parentOf(removed));
}
setColor(sibling, colorOf(parentOf(removed)));
setColor(parentOf(removed), BLACK);
setColor(rightOf(sibling), BLACK);
rotateLeft(parentOf(removed));
removed = root;
}
} else {
TreeNode<K, V> sibling = leftOf(parentOf(removed));
if (colorOf(sibling) == RED) {
setColor(sibling, BLACK);
setColor(parentOf(removed), RED);
rotateRight(parentOf(removed));
sibling = leftOf(parentOf(removed));
}
if (colorOf(leftOf(sibling)) == BLACK && colorOf(parentOf(sibling)) == BLACK) {
setColor(sibling, RED);
removed = parentOf(removed);
} else {
if (colorOf(leftOf(sibling)) == BLACK) {
setColor(sibling, RED);
setColor(rightOf(sibling), BLACK);
rotateLeft(sibling);
sibling = leftOf(parentOf(removed));
}
setColor(sibling, colorOf(parentOf(removed)));
setColor(parentOf(removed), BLACK);
setColor(leftOf(sibling), BLACK);
rotateRight(parentOf(removed));
removed = root;
}
}
}
setColor(removed, BLACK);
}
}
Java 实现红黑树
最新推荐文章于 2024-08-07 09:27:53 发布