use std::cell::RefCell;
use std::rc::Rc;
use std::fmt::Debug;
use crate::binary_tree_ops::{BinaryTreeOps, TreeNode, NodeColor, NodeOps};
pub struct RedBlackTreeImpl<T: Ord + Clone + Debug> {
root: Option<Rc<RefCell<TreeNode<T>>>>,
}
impl<T: Ord + Clone + Debug> RedBlackTreeImpl<T> {
pub fn new() -> Self {
Self { root: None }
}
pub fn insert(&mut self, key: T) {
let new_node = TreeNode::new(key);
new_node.borrow_mut().set_color(Some(NodeColor::Red));
if let Some(root) = self.root_node() {
self.rb_insert_recursive(root, new_node.clone());
} else {
self.set_root_node(Some(new_node.clone()));
}
self.fix_insert(new_node);
}
fn rb_insert_recursive(&self, current: Rc<RefCell<TreeNode<T>>>, new_node: Rc<RefCell<TreeNode<T>>>) {
let new_key = new_node.borrow().key.clone();
let mut curr_borrow = current.borrow_mut();
if new_key < curr_borrow.key {
if let Some(left_child) = curr_borrow.left.clone() {
drop(curr_borrow);
self.rb_insert_recursive(left_child, new_node);
} else {
new_node.borrow_mut().parent = Some(current.clone());
curr_borrow.left = Some(new_node);
}
} else {
if let Some(right_child) = curr_borrow.right.clone() {
drop(curr_borrow);
self.rb_insert_recursive(right_child, new_node);
} else {
new_node.borrow_mut().parent = Some(current.clone());
curr_borrow.right = Some(new_node);
}
}
}
fn rotate_left(&mut self, node: Rc<RefCell<TreeNode<T>>>) {
let right = node.borrow_mut().right.take();
if let Some(r) = right {
let parent = node.borrow().parent.clone();
r.borrow_mut().parent = parent.clone();
if let Some(p) = parent {
let is_left = p.borrow().left.as_ref().map_or(false, |l| Rc::ptr_eq(l, &node));
if is_left {
p.borrow_mut().left = Some(r.clone());
} else {
p.borrow_mut().right = Some(r.clone());
}
} else {
self.set_root_node(Some(r.clone()));
}
{
let mut nb = node.borrow_mut();
let left_sub = r.borrow_mut().left.take();
nb.right = left_sub;
if let Some(l) = &nb.right {
l.borrow_mut().parent = Some(node.clone());
}
}
r.borrow_mut().left = Some(node.clone());
node.borrow_mut().parent = Some(r);
}
}
fn rotate_right(&mut self, node: Rc<RefCell<TreeNode<T>>>) {
let left = node.borrow_mut().left.take();
if let Some(l) = left {
let parent = node.borrow().parent.clone();
l.borrow_mut().parent = parent.clone();
if let Some(p) = parent {
let is_right = p.borrow().right.as_ref().map_or(false, |r| Rc::ptr_eq(r, &node));
if is_right {
p.borrow_mut().right = Some(l.clone());
} else {
p.borrow_mut().left = Some(l.clone());
}
} else {
self.set_root_node(Some(l.clone()));
}
{
let mut nb = node.borrow_mut();
let right_sub = l.borrow_mut().right.take();
nb.left = right_sub;
if let Some(r) = &nb.left {
r.borrow_mut().parent = Some(node.clone());
}
}
l.borrow_mut().right = Some(node.clone());
node.borrow_mut().parent = Some(l);
}
}
fn replace_node(&mut self, old: Rc<RefCell<TreeNode<T>>>, new: Option<Rc<RefCell<TreeNode<T>>>>) {
let parent = old.borrow().parent.clone();
if let Some(p) = parent {
let is_left = p.borrow().left.as_ref().map_or(false, |l| Rc::ptr_eq(l, &old));
if is_left {
p.borrow_mut().left = new.clone();
} else {
p.borrow_mut().right = new.clone();
}
if let Some(nw) = new {
nw.borrow_mut().parent = Some(p);
}
} else {
self.set_root_node(new);
}
}
fn node_color(n: &Option<Rc<RefCell<TreeNode<T>>>>) -> NodeColor {
if let Some(node) = n {
node.borrow().color().unwrap_or(NodeColor::Black)
} else {
NodeColor::Black
}
}
fn fix_insert(&mut self, node: Rc<RefCell<TreeNode<T>>>) {
let mut current = node;
loop {
let parent_opt = {
let cb = current.borrow();
cb.parent.clone()
};
if parent_opt.is_none() {
break;
}
let mut parent = parent_opt.unwrap();
if parent.borrow().color() == Some(NodeColor::Black) {
break;
}
let grandparent_opt = {
let pb = parent.borrow();
pb.parent.clone()
};
if grandparent_opt.is_none() {
break;
}
let mut grandparent = grandparent_opt.unwrap();
let parent_is_left = {
let g = grandparent.borrow();
g.left.as_ref().map_or(false, |l| Rc::ptr_eq(l, &parent))
};
let uncle = if parent_is_left {
grandparent.borrow().right.clone()
} else {
grandparent.borrow().left.clone()
};
// Case 1: Uncle is Red - Recolor and move up
if let Some(u) = uncle {
if u.borrow().color() == Some(NodeColor::Red) {
parent.borrow_mut().set_color(Some(NodeColor::Black));
u.borrow_mut().set_color(Some(NodeColor::Black));
grandparent.borrow_mut().set_color(Some(NodeColor::Red));
current = grandparent;
continue;
}
}
// Case 2/3: Uncle is Black - May need rotations
if parent_is_left {
// Left-Right case
let parent_right = {
let pb = parent.borrow();
pb.right.clone()
};
if parent_right.as_ref().map_or(false, |r| Rc::ptr_eq(r, ¤t)) {
// Rotate left on parent to resolve Left-Right
current = parent.clone();
self.rotate_left(current.clone());
// Re-identify parent and grandparent after rotation
let new_parent_rc = {
let c_ref = current.borrow();
c_ref.parent.clone()
};
if let Some(new_parent_rc) = new_parent_rc {
let new_grandparent_rc_opt = {
let p_ref = new_parent_rc.borrow();
p_ref.parent.clone()
};
if let Some(new_grandparent_rc) = new_grandparent_rc_opt {
parent = new_parent_rc;
grandparent = new_grandparent_rc;
}
}
}
// Now handle the Left-Left case
parent.borrow_mut().set_color(Some(NodeColor::Black));
grandparent.borrow_mut().set_color(Some(NodeColor::Red));
self.rotate_right(grandparent.clone());
} else {
// Right-Left case
let parent_left = {
let pb = parent.borrow();
pb.left.clone()
};
if parent_left.as_ref().map_or(false, |l| Rc::ptr_eq(l, ¤t)) {
// Rotate right on parent to resolve Right-Left
current = parent.clone();
self.rotate_right(current.clone());
// Re-identify parent and grandparent after rotation
let new_parent_rc = {
let c_ref = current.borrow();
c_ref.parent.clone()
};
if let Some(new_parent_rc) = new_parent_rc {
let new_grandparent_rc_opt = {
let p_ref = new_parent_rc.borrow();
p_ref.parent.clone()
};
if let Some(new_grandparent_rc) = new_grandparent_rc_opt {
parent = new_parent_rc;
grandparent = new_grandparent_rc;
}
}
}
// Now handle the Right-Right case
parent.borrow_mut().set_color(Some(NodeColor::Black));
grandparent.borrow_mut().set_color(Some(NodeColor::Red));
self.rotate_left(grandparent.clone());
}
}
// Ensure root is always black
if let Some(r) = self.root_node() {
r.borrow_mut().set_color(Some(NodeColor::Black));
}
}
pub fn delete(&mut self, key: T) {
if let Some(node) = self.find_node(&key) {
self.rb_delete_node(node);
}
}
fn rb_delete_node(&mut self, node: Rc<RefCell<TreeNode<T>>>) {
// Standard RB-Tree deletion procedure:
let mut y = node.clone();
let y_original_color = y.borrow().color().unwrap_or(NodeColor::Black);
let (left, right) = {
let nb = node.borrow();
(nb.left.clone(), nb.right.clone())
};
#[allow(unused_mut)]
let mut x: Option<Rc<RefCell<TreeNode<T>>>>;
if left.is_none() {
x = right.clone();
self.replace_node(node.clone(), right);
} else if right.is_none() {
x = left.clone();
self.replace_node(node.clone(), left);
} else {
let successor = self.find_min(node.borrow().right.clone().unwrap());
let successor_color = successor.borrow().color().unwrap_or(NodeColor::Black);
y = successor.clone();
let y_right = y.borrow().right.clone();
x = y_right.clone();
if Rc::ptr_eq(&y.borrow().parent.as_ref().unwrap(), &node) {
if let Some(xx) = &x {
xx.borrow_mut().parent = Some(y.clone());
}
} else {
self.replace_node(y.clone(), y_right);
y.borrow_mut().right = node.borrow().right.clone();
if let Some(nr) = &y.borrow().right {
nr.borrow_mut().parent = Some(y.clone());
}
}
self.replace_node(node.clone(), Some(y.clone()));
y.borrow_mut().left = node.borrow().left.clone();
if let Some(nl) = &y.borrow().left {
nl.borrow_mut().parent = Some(y.clone());
}
y.borrow_mut().set_color(node.borrow().color());
if successor_color == NodeColor::Black {
self.fix_delete(x);
}
return;
}
if y_original_color == NodeColor::Black {
self.fix_delete(x);
}
}
fn fix_delete(&mut self, mut x: Option<Rc<RefCell<TreeNode<T>>>>) {
while x.as_ref().map_or(false, |xx| !Rc::ptr_eq(xx, &self.root_node().unwrap()))
&& Self::node_color(&x) == NodeColor::Black
{
let x_parent = x.as_ref().and_then(|xx| xx.borrow().parent.clone());
if let Some(p) = x_parent {
let x_is_left = p.borrow().left.as_ref().map_or(false, |l| {
if let Some(xx) = &x { Rc::ptr_eq(l, xx) } else { false }
});
let mut w = if x_is_left {
p.borrow().right.clone()
} else {
p.borrow().left.clone()
};
if Self::node_color(&w) == NodeColor::Red {
// Case 1
if let Some(ww) = w.clone() {
ww.borrow_mut().set_color(Some(NodeColor::Black));
}
p.borrow_mut().set_color(Some(NodeColor::Red));
if x_is_left {
self.rotate_left(p.clone());
} else {
self.rotate_right(p.clone());
}
w = if x_is_left {
p.borrow().right.clone()
} else {
p.borrow().left.clone()
};
}
// Case 2
if Self::node_color(&w.as_ref().and_then(|ww| ww.borrow().left.clone())) == NodeColor::Black &&
Self::node_color(&w.as_ref().and_then(|ww| ww.borrow().right.clone())) == NodeColor::Black
{
if let Some(ww) = w.clone() {
ww.borrow_mut().set_color(Some(NodeColor::Red));
}
x = p.clone().into();
} else {
// Case 3/4
if x_is_left {
if Self::node_color(&w.as_ref().and_then(|ww| ww.borrow().right.clone())) == NodeColor::Black {
if let Some(wl) = w.as_ref().and_then(|ww| ww.borrow().left.clone()) {
wl.borrow_mut().set_color(Some(NodeColor::Black));
}
if let Some(ww) = w.clone() {
ww.borrow_mut().set_color(Some(NodeColor::Red));
}
if let Some(ww) = w.clone() {
self.rotate_right(ww);
}
w = p.borrow().right.clone();
}
if let Some(ww) = w.clone() {
ww.borrow_mut().set_color(p.borrow().color());
}
p.borrow_mut().set_color(Some(NodeColor::Black));
if let Some(wr) = w.as_ref().and_then(|ww| ww.borrow().right.clone()) {
wr.borrow_mut().set_color(Some(NodeColor::Black));
}
self.rotate_left(p.clone());
x = self.root_node();
} else {
if Self::node_color(&w.as_ref().and_then(|ww| ww.borrow().left.clone())) == NodeColor::Black {
if let Some(wr) = w.as_ref().and_then(|ww| ww.borrow().right.clone()) {
wr.borrow_mut().set_color(Some(NodeColor::Black));
}
if let Some(ww) = w.clone() {
ww.borrow_mut().set_color(Some(NodeColor::Red));
}
if let Some(ww) = w.clone() {
self.rotate_left(ww);
}
w = p.borrow().left.clone();
}
if let Some(ww) = w.clone() {
ww.borrow_mut().set_color(p.borrow().color());
}
p.borrow_mut().set_color(Some(NodeColor::Black));
if let Some(wl) = w.as_ref().and_then(|ww| ww.borrow().left.clone()) {
wl.borrow_mut().set_color(Some(NodeColor::Black));
}
self.rotate_right(p.clone());
x = self.root_node();
}
}
} else {
break;
}
}
if let Some(xx) = x {
xx.borrow_mut().set_color(Some(NodeColor::Black));
}
if let Some(r) = self.root_node() {
r.borrow_mut().set_color(Some(NodeColor::Black));
}
}
fn find_min(&self, start_node: Rc<RefCell<TreeNode<T>>>) -> Rc<RefCell<TreeNode<T>>> {
let mut node = start_node;
while let Some(left) = {
let nb = node.borrow();
nb.left.clone()
} {
node = left;
}
node
}
}
impl<T: Ord + Clone + Debug> BinaryTreeOps<T> for RedBlackTreeImpl<T> {
fn root_node(&self) -> Option<Rc<RefCell<TreeNode<T>>>> {
self.root.clone()
}
fn set_root_node(&mut self, root: Option<Rc<RefCell<TreeNode<T>>>>) {
self.root = root;
if let Some(r) = &self.root {
r.borrow_mut().parent = None;
r.borrow_mut().set_color(Some(NodeColor::Black));
}
}
}