1use crate::node::{Node, NodeColor};
2use std::{
3 cell::RefCell,
4 fmt::{Debug, Formatter},
5 rc::Rc,
6};
7
8#[cfg(test)]
9mod tree_tests;
10
11pub struct Tree<T> {
12 root: Rc<RefCell<Node<T>>>,
13 sentinel: Rc<RefCell<Node<T>>>,
14 length: usize,
15}
16
17impl<T: PartialOrd + Clone + PartialEq + Debug + Default> Tree<T> {
19 pub fn new() -> Tree<T> {
20 let sentinel = Rc::new(RefCell::new(Node::new_sentinel()));
21 Self {
22 root: sentinel.clone(),
23 sentinel,
24 length: 0,
25 }
26 }
27
28 pub fn insert(&mut self, key: T) {
29 let mut z = Node::new(key);
30 let mut x = self.root.clone();
31 let mut y = self.sentinel.clone();
32
33 while !x.borrow().is_nil() {
34 y = x.clone();
35 if z.key < x.borrow().key {
36 let x_tmp = x.borrow().left().clone();
37 x = x_tmp
38 } else {
39 let x_tmp = x.borrow().right().clone();
40 x = x_tmp;
41 }
42 }
43 z.set_parent(y.clone());
44 let z = Rc::new(RefCell::new(z));
46
47 if y.borrow().is_nil() {
48 self.root = z.clone();
49 } else if z.borrow().key < y.borrow().key {
50 y.borrow_mut().set_left_child(z.clone());
51 } else {
52 y.borrow_mut().set_right_child(z.clone());
53 }
54
55 z.borrow_mut().set_left_child(self.sentinel.clone());
56 z.borrow_mut().set_right_child(self.sentinel.clone());
57 z.borrow_mut().color = NodeColor::Red;
58 self.insert_fix_up(z);
59 self.length += 1;
60 }
61
62 fn insert_fix_up(&mut self, mut z: Rc<RefCell<Node<T>>>) {
63 while z.borrow().parent().borrow().color == NodeColor::Red {
64 if z.borrow().parent() == z.borrow().parent().borrow().parent().borrow().left() {
65 let y = z
66 .borrow()
67 .parent()
68 .borrow()
69 .parent()
70 .borrow()
71 .right()
72 .clone();
73 if y.borrow().color == NodeColor::Red {
75 z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
76 y.borrow_mut().color = NodeColor::Black;
77 z.borrow_mut()
78 .parent_mut()
79 .borrow_mut()
80 .parent_mut()
81 .borrow_mut()
82 .color = NodeColor::Red;
83 let z_tmp = z.borrow().parent().borrow().parent().clone();
84 z = z_tmp;
85 } else {
86 if &z == z.borrow().parent().borrow().right() {
88 let z_tmp = z.borrow().parent().clone();
89 z = z_tmp;
90 self.left_rotate(z.clone());
91 }
92 z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
94 z.borrow_mut()
95 .parent_mut()
96 .borrow_mut()
97 .parent_mut()
98 .borrow_mut()
99 .color = NodeColor::Red;
100 let x = z.borrow().parent().borrow().parent().clone();
101 self.right_rotate(x);
102 }
103 } else {
104 let y = z
105 .borrow()
106 .parent()
107 .borrow()
108 .parent()
109 .borrow()
110 .left()
111 .clone();
112 if y.borrow().color == NodeColor::Red {
114 z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
115 y.borrow_mut().color = NodeColor::Black;
116 z.borrow_mut()
117 .parent_mut()
118 .borrow_mut()
119 .parent_mut()
120 .borrow_mut()
121 .color = NodeColor::Red;
122 let z_tmp = z.borrow().parent().borrow().parent().clone();
123 z = z_tmp;
124 } else {
125 if &z == z.borrow().parent().borrow().left() {
127 let z_tmp = z.borrow().parent().clone();
128 z = z_tmp;
129 self.right_rotate(z.clone());
130 }
131 z.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
133 z.borrow_mut()
134 .parent_mut()
135 .borrow_mut()
136 .parent_mut()
137 .borrow_mut()
138 .color = NodeColor::Red;
139 let x = z.borrow().parent().borrow().parent().clone();
140 self.left_rotate(x);
141 }
142 }
143 }
144 self.root.borrow_mut().color = NodeColor::Black;
145 }
146
147 fn left_rotate(&mut self, x: Rc<RefCell<Node<T>>>) {
148 if !x.borrow().right().borrow().is_nil() {
150 let y = x.borrow().right().clone();
152 x.borrow_mut().set_right_child(y.borrow().left().clone());
153 if !y.borrow().left().borrow().is_nil() {
154 y.borrow_mut().left_mut().borrow_mut().set_parent(x.clone());
155 }
156 y.borrow_mut().set_parent(x.borrow().parent().clone());
157 if x.borrow().parent().borrow().is_nil() {
158 self.root = y.clone();
159 } else if &x == x.borrow().parent().borrow().left() {
160 x.borrow_mut()
161 .parent_mut()
162 .borrow_mut()
163 .set_left_child(y.clone());
164 } else {
165 x.borrow_mut()
166 .parent_mut()
167 .borrow_mut()
168 .set_right_child(y.clone());
169 }
170
171 y.borrow_mut().set_left_child(x.clone());
172 x.borrow_mut().set_parent(y);
173 } else {
174 panic!(
175 "Invariant violated. The right child of {:?} must not be T.nil.",
176 x
177 );
178 }
179 }
180
181 fn right_rotate(&mut self, y: Rc<RefCell<Node<T>>>) {
182 if !y.borrow().left().borrow().is_nil() {
184 let x = y.borrow().left().clone();
186 y.borrow_mut().set_left_child(x.borrow().right().clone());
187 if !x.borrow().right().borrow().is_nil() {
188 x.borrow_mut()
189 .right_mut()
190 .borrow_mut()
191 .set_parent(y.clone());
192 }
193 x.borrow_mut().set_parent(y.borrow().parent().clone());
194 if y.borrow().parent().borrow().is_nil() {
195 self.root = x.clone();
196 } else if &y.clone() == y.borrow().parent().borrow().right() {
197 y.borrow_mut()
198 .parent_mut()
199 .borrow_mut()
200 .set_right_child(x.clone());
201 } else {
202 y.borrow_mut()
203 .parent_mut()
204 .borrow_mut()
205 .set_left_child(x.clone());
206 }
207 x.borrow_mut().set_right_child(y.clone());
208 y.borrow_mut().set_parent(x);
209 } else {
210 panic!(
211 "Invariant violated. The left child of {:?} must not be T.nil.",
212 y
213 );
214 }
215 }
216
217 pub fn delete(&mut self, key: T) {
218 let node = self.search(key);
219 if node.is_some() {
220 self.delete_node(node.as_ref().unwrap().clone())
221 }
222 self.length -= 1;
223 }
224
225 fn delete_node(&mut self, z: Rc<RefCell<Node<T>>>) {
226 let mut y = z.clone();
227 let mut y_color = y.borrow().color.clone();
228 let x;
229 if z.borrow().left().borrow().is_nil() {
230 x = z.borrow().right().clone();
231 let u = z.clone();
232 let v = z.borrow().right().clone();
233 self.transplant(u, v);
234 } else if z.borrow().right().borrow().is_nil() {
235 x = z.borrow().left().clone();
236 let u = z.clone();
237 let v = z.borrow().left().clone();
238 self.transplant(u, v);
239 } else {
240 y = self
241 .minimum_node(z.borrow().right().clone())
242 .expect("Expected this to be set");
243 y_color = y.borrow().color.clone();
244 x = y.borrow().right().clone();
245 if &y != z.borrow().right() {
246 let u = y.clone();
247 let v = y.borrow().right().clone();
248 self.transplant(u, v);
249 y.borrow_mut().set_right_child(z.borrow().right().clone());
250 y.borrow_mut()
251 .right_mut()
252 .borrow_mut()
253 .set_parent(y.clone());
254 } else {
255 x.borrow_mut().set_parent(y.clone());
256 }
257 let u = z.clone();
258 let v = y.clone();
259 self.transplant(u, v);
260 y.borrow_mut().set_left_child(z.borrow().left().clone());
261 y.borrow_mut().left_mut().borrow_mut().set_parent(y.clone());
262 y.borrow_mut().color = z.borrow().color.clone();
263 }
264
265 if y_color == NodeColor::Black {
266 self.delete_fix_up(x);
267 }
268 }
269
270 fn transplant(&mut self, u: Rc<RefCell<Node<T>>>, v: Rc<RefCell<Node<T>>>) {
271 if u.borrow().parent().borrow().is_nil() {
272 self.root = v.clone();
273 } else if &u == u.borrow().parent().borrow().left() {
274 u.borrow_mut()
275 .parent_mut()
276 .borrow_mut()
277 .set_left_child(v.clone());
278 } else {
279 u.borrow_mut()
280 .parent_mut()
281 .borrow_mut()
282 .set_right_child(v.clone());
283 }
284 v.borrow_mut().set_parent(u.borrow().parent().clone());
285 }
286
287 fn delete_fix_up(&mut self, mut x: Rc<RefCell<Node<T>>>) {
288 while x != self.root && x.borrow().color == NodeColor::Black {
289 if &x == x.borrow().parent().borrow().left() {
290 let mut w = x.borrow().parent().borrow().right().clone();
291 if w.borrow().color == NodeColor::Red {
293 w.borrow_mut().color = NodeColor::Black;
294 x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Red;
295 self.left_rotate(x.borrow().parent().clone());
296 w = x.borrow().parent().borrow().right().clone();
297 }
298
299 if w.borrow().left().borrow().color == NodeColor::Black
300 && w.borrow().right().borrow().color == NodeColor::Black
301 {
302 w.borrow_mut().color = NodeColor::Red;
303 let x_tmp = x.borrow().parent().clone();
304 x = x_tmp;
305 } else {
306 if w.borrow_mut().right().borrow().color == NodeColor::Black {
308 w.borrow_mut().left().borrow_mut().color = NodeColor::Black;
309 w.borrow_mut().color = NodeColor::Red;
310 self.right_rotate(w.clone());
311 w = x.borrow().parent().borrow().right().clone();
312 }
313 w.borrow_mut().color = x.borrow().parent().borrow().color.clone();
315 x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
316 w.borrow_mut().right_mut().borrow_mut().color = NodeColor::Black;
317 self.left_rotate(x.borrow().parent().clone());
318 x = self.root.clone();
319 }
320 } else {
321 let mut w = x.borrow().parent().borrow().left().clone();
322 if w.borrow_mut().color == NodeColor::Red {
324 w.borrow_mut().color = NodeColor::Black;
325 x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Red;
326 self.right_rotate(x.borrow().parent().clone());
327 w = x.borrow().parent().borrow().left().clone();
328 }
329 if w.borrow().right().borrow().color == NodeColor::Black
331 && w.borrow().left().borrow().color == NodeColor::Black
332 {
333 w.borrow_mut().color = NodeColor::Red;
334 let x_tmp = x.borrow().parent().clone();
335 x = x_tmp;
336 } else {
337 if w.borrow().left().borrow().color == NodeColor::Black {
339 w.borrow_mut().right_mut().borrow_mut().color = NodeColor::Black;
340 w.borrow_mut().color = NodeColor::Red;
341 self.left_rotate(w.clone());
342 w = x.borrow_mut().parent().borrow().left().clone();
343 }
344 w.borrow_mut().color = x.borrow().parent().borrow().color.clone();
346 x.borrow_mut().parent_mut().borrow_mut().color = NodeColor::Black;
347 w.borrow_mut().left_mut().borrow_mut().color = NodeColor::Black;
348 self.right_rotate(x.borrow().parent().clone());
349 x = self.root.clone();
350 }
351 }
352 }
353 x.borrow_mut().color = NodeColor::Black;
354 }
355
356 pub fn contains_key(&self, key: T) -> bool {
357 self.search(key).is_some()
358 }
359
360 fn search(&self, key: T) -> Option<Rc<RefCell<Node<T>>>> {
361 let mut node = self.root.clone();
362 while !node.borrow().is_nil() {
363 let node_key = node.borrow().key.clone();
364 if node_key == key {
365 return Some(node);
366 } else if key < node_key {
367 let node_tmp = node.borrow().left().clone();
368 node = node_tmp;
369 } else {
370 let node_tmp = node.borrow().right().clone();
371 node = node_tmp;
372 }
373 }
374 None
375 }
376
377 pub fn is_empty(&self) -> bool {
378 self.length == 0
379 }
380
381 pub fn len(&self) -> usize {
382 self.length
383 }
384
385 pub fn clear(&mut self) {
386 self.root = self.sentinel.clone();
387 self.length = 0;
388 }
389
390 pub fn minimum(&self) -> Option<T> {
391 if self.root.borrow().is_nil() {
392 return None;
393 }
394
395 self.minimum_node(self.root.clone())
396 .and_then(|node| Some(node.borrow().key.clone()))
397 }
398 fn minimum_node(&self, node: Rc<RefCell<Node<T>>>) -> Option<Rc<RefCell<Node<T>>>> {
399 if node.borrow().left().borrow().is_nil() {
400 return Some(node);
401 }
402
403 let mut x = node.clone();
404 while !x.borrow().left().borrow().is_nil() {
405 let x_tmp = x.borrow().left().clone();
406 x = x_tmp;
407 }
408 Some(x)
409 }
410
411 pub fn maximum(&self) -> Option<T> {
412 if self.root.borrow().is_nil() {
413 return None;
414 }
415
416 self.maximum_node(self.root.clone())
417 .and_then(|node| Some(node.borrow().key.clone()))
418 }
419 fn maximum_node(&self, node: Rc<RefCell<Node<T>>>) -> Option<Rc<RefCell<Node<T>>>> {
420 if node.borrow().right().borrow().is_nil() {
421 return Some(node);
422 }
423
424 let mut x = node.clone();
425 while !x.borrow().right().borrow().is_nil() {
426 let x_tmp = x.borrow().right().clone();
427 x = x_tmp;
428 }
429 Some(x)
430 }
431}
432
433impl<T: Debug> Debug for Tree<T> {
434 fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
435 write!(f, "{:?}", self.root)
436 }
437}
438
439fn tree_equality_dfs<T: PartialEq>(
443 me: Option<Rc<RefCell<Node<T>>>>,
444 other: Option<Rc<RefCell<Node<T>>>>,
445) -> bool {
446 match (me.clone(), other.clone()) {
448 (None, None) => panic!("Invariant violated. Node's cannot be None"),
449 (Some(_), Some(_)) => (),
450 _ => return false,
451 }
452
453 if me.as_ref().unwrap().borrow().color != other.as_ref().unwrap().borrow().color {
454 return false;
455 }
456
457 if me.as_ref().unwrap().borrow().key != other.as_ref().unwrap().borrow().key {
458 return false;
459 }
460
461 match (
462 me.as_ref().unwrap().borrow().is_sentinel,
463 other.as_ref().unwrap().borrow().is_sentinel,
464 ) {
465 (true, true) => return true,
466 (false, false) => (),
467 _ => return false,
468 }
469
470 let left_subtree = tree_equality_dfs(
471 me.as_ref().unwrap().borrow().left.clone(),
472 other.as_ref().unwrap().borrow().left.clone(),
473 );
474 if !left_subtree {
475 return false;
476 }
477
478 let right_subtree = tree_equality_dfs(
479 me.as_ref().unwrap().borrow().right.clone(),
480 other.as_ref().unwrap().borrow().right.clone(),
481 );
482 if !right_subtree {
483 return false;
484 }
485
486 true
487}
488
489impl<T: PartialEq> PartialEq<Self> for Tree<T> {
490 fn eq(&self, other: &Self) -> bool {
491 tree_equality_dfs(Some(self.root.clone()), Some(other.root.clone()))
492 }
493}