rudac/tree/
rb.rs

1use std::collections::VecDeque;
2
3const RED: bool = true;
4const BLACK: bool = false;
5
6struct Node<K: std::cmp::Ord, V> {
7    key: Option<K>,
8    value: Option<V>,
9    color: bool,
10    size: usize,
11    left_child: Option<Box<Node<K, V>>>,
12    right_child: Option<Box<Node<K, V>>>,
13}
14
15impl<K: std::cmp::Ord, V> Node<K, V> {
16    fn init(key: K, value: V, color: bool, size: usize) -> Node<K, V> {
17        Node {
18            key: Some(key),
19            value: Some(value),
20            color: color,
21            size: size,
22            left_child: None,
23            right_child: None,
24        }
25    }
26
27    fn key(&self) -> &K {
28        &self.key.as_ref().unwrap()
29    }
30    fn key_mut(&mut self) -> &mut Option<K> {
31        &mut self.key
32    }
33
34    fn get_key(&mut self) -> K {
35        self.key.take().unwrap()
36    }
37
38    fn value(&self) -> &V {
39        &self.value.as_ref().unwrap()
40    }
41
42    fn value_mut(&mut self) -> &mut Option<V> {
43        &mut self.value
44    }
45
46    fn get_value(&mut self) -> V {
47        self.value.take().unwrap()
48    }
49
50    fn left_child(&self) -> &Box<Node<K, V>> {
51        self.left_child.as_ref().unwrap()
52    }
53
54    fn right_child(&self) -> &Box<Node<K, V>> {
55        self.right_child.as_ref().unwrap()
56    }
57
58    fn is_red(node: &Option<Box<Node<K, V>>>) -> bool {
59        if node.is_none() {
60            return false;
61        }
62        node.as_ref().unwrap().color == RED
63    }
64
65    fn size(node: &Option<Box<Node<K, V>>>) -> usize {
66        if node.is_none() {
67            return 0;
68        }
69        node.as_ref().unwrap().size
70    }
71
72    fn update_size(&mut self) {
73        self.size = Node::size(&self.left_child) + Node::size(&self.right_child) + 1;
74    }
75}
76
77/// A Red Black tree is a self-balancing binary search tree.
78/// Red Black Trees provide faster insertion and removal operations than AVL trees
79///
80/// # Examples
81/// ```
82/// use rudac::tree::RedBlack;
83///
84/// // initialize a Red Black tree with keys of type usize and values of type String
85/// let mut rb_tree = RedBlack::<usize, String>::init();
86///
87/// // insert items into tree
88/// rb_tree.insert(1, String::from("rudac"));
89/// rb_tree.insert(2, String::from("is"));
90/// rb_tree.insert(3, String::from("awesome"));
91/// rb_tree.insert(4, String::from("!"));
92///
93/// // lookup for items
94/// assert_eq!(*rb_tree.get(&1).unwrap(), String::from("rudac"));
95/// assert_eq!(*rb_tree.get(&2).unwrap(), String::from("is"));
96/// assert_eq!(*rb_tree.get(&3).unwrap(), String::from("awesome"));
97/// assert_eq!(*rb_tree.get(&4).unwrap(), String::from("!"));
98///
99/// // delete items from tree
100/// rb_tree.delete(&4);
101/// assert_eq!(rb_tree.get(&4), None);
102/// ```
103pub struct RedBlack<K: std::cmp::Ord, V> {
104    root: Option<Box<Node<K, V>>>,
105}
106
107impl<K: std::cmp::Ord, V> RedBlack<K, V> {
108    /// Initializes an empty Red Black tree
109    ///
110    /// # Examples
111    /// ```
112    /// use rudac::tree::RedBlack;
113    ///
114    /// let usize_to_string = RedBlack::<usize, String>::init();
115    ///
116    /// let string_to_usize = RedBlack::<String, usize>::init();
117    ///
118    /// let string_to_string = RedBlack::<String, String>::init();
119    /// ```
120    pub fn init() -> RedBlack<K, V> {
121        RedBlack { root: None }
122    }
123
124    /// Returns total number of nodes in the tree
125    ///
126    /// # Examples
127    /// ```
128    /// use rudac::tree::RedBlack;
129    ///
130    /// let mut rb_tree = RedBlack::<usize,usize>::init();
131    /// assert_eq!(rb_tree.size(), 0);
132    ///
133    /// rb_tree.insert(1,1);
134    /// rb_tree.insert(2,4);
135    /// rb_tree.insert(3,8);
136    /// assert_eq!(rb_tree.size(), 3);
137    /// ```
138    pub fn size(&self) -> usize {
139        Node::size(&self.root)
140    }
141
142    /// Returns `true` if tree is empty and `false` otherwise
143    ///
144    /// # Examples
145    /// ```
146    /// use rudac::tree::RedBlack;
147    ///
148    /// let mut rb_tree = RedBlack::<usize,usize>::init();
149    /// assert_eq!(rb_tree.is_empty(), true);
150    ///
151    /// rb_tree.insert(1,1);
152    /// assert_eq!(rb_tree.is_empty(), false);
153    ///
154    /// rb_tree.delete(&1);
155    /// assert_eq!(rb_tree.is_empty(), true);
156    /// ```
157    pub fn is_empty(&self) -> bool {
158        self.root.is_none()
159    }
160
161    /// Returns a reference to value associated with specified `key` in tree, `None` otherwise
162    /// # Arguments
163    /// * `key`: key to be searched in the tree
164    ///
165    /// # Examples
166    /// ```
167    /// use rudac::tree::RedBlack;
168    ///
169    /// let mut rb_tree = RedBlack::<usize,usize>::init();
170    ///
171    /// rb_tree.insert(1,10);
172    /// assert_eq!(*rb_tree.get(&1).unwrap(), 10);
173    ///
174    /// rb_tree.delete(&1);
175    /// assert_eq!(rb_tree.get(&1), None);
176    /// ```
177    pub fn get(&self, key: &K) -> Option<&V> {
178        RedBlack::_get(&self.root, key)
179    }
180
181    fn _get<'a>(mut node: &'a Option<Box<Node<K, V>>>, key: &K) -> Option<&'a V> {
182        while !node.is_none() {
183            let node_ref = node.as_ref().unwrap();
184            if key < node_ref.key() {
185                node = &node_ref.left_child
186            } else if key > node_ref.key() {
187                node = &node_ref.right_child
188            } else {
189                return Some(node_ref.value());
190            }
191        }
192        None
193    }
194
195    /// Returns `true` if tree contains the specified `key`, false otherwise
196    ///
197    /// # Arguments
198    /// * `key`: key to be searched in the tree
199    ///
200    /// # Examples
201    /// ```
202    /// use rudac::tree::RedBlack;
203    ///
204    /// let mut rb_tree = RedBlack::<usize,usize>::init();
205    ///
206    /// rb_tree.insert(1,10);
207    /// assert_eq!(rb_tree.contains(&1), true);
208    ///
209    /// rb_tree.delete(&1);
210    /// assert_eq!(rb_tree.contains(&1), false);
211    /// ```
212    pub fn contains(&self, key: &K) -> bool {
213        !self.get(key).is_none()
214    }
215
216    /// Insert a node which contains the specified `key` and `value` into the tree.
217    /// if `key` already exists, this method will replace `value` as the new value of the node
218    ///
219    /// # Arguments
220    /// * `key`: key of the new node
221    /// * `value`: value associated with the `key`
222    ///
223    /// # Examples
224    /// ```
225    /// use rudac::tree::RedBlack;
226    ///
227    /// let mut rb_tree = RedBlack::<usize,usize>::init();
228    ///
229    /// rb_tree.insert(1,10);
230    /// rb_tree.insert(2,20);
231    /// rb_tree.insert(3,30);
232    /// rb_tree.insert(4,40);
233    /// assert_eq!(*rb_tree.get(&1).unwrap(), 10);
234    ///
235    /// rb_tree.insert(1,11);
236    /// assert_eq!(*rb_tree.get(&1).unwrap(), 11);
237    /// ```
238    pub fn insert(&mut self, key: K, value: V) {
239        let mut root = RedBlack::_insert(self.root.take(), key, value).unwrap();
240
241        root.color = BLACK;
242
243        self.root = Some(root);
244    }
245
246    fn _insert(node: Option<Box<Node<K, V>>>, key: K, value: V) -> Option<Box<Node<K, V>>> {
247        if node.is_none() {
248            return Some(Box::new(Node::init(key, value, RED, 1)));
249        }
250
251        let mut node_ref = node.unwrap();
252
253        if key < *node_ref.key() {
254            node_ref.left_child = RedBlack::_insert(node_ref.left_child, key, value);
255        } else if key > *node_ref.key() {
256            node_ref.right_child = RedBlack::_insert(node_ref.right_child, key, value);
257        } else {
258            node_ref.value = Some(value);
259        }
260
261        // balance the tree
262        if Node::is_red(&node_ref.right_child) && !Node::is_red(&node_ref.left_child) {
263            node_ref = RedBlack::rotate_left(node_ref);
264        }
265        if Node::is_red(&node_ref.left_child) && Node::is_red(&node_ref.left_child().left_child) {
266            node_ref = RedBlack::rotate_right(node_ref);
267        }
268        if Node::is_red(&node_ref.left_child) && Node::is_red(&node_ref.right_child) {
269            RedBlack::flip_colors(&mut node_ref);
270        }
271
272        node_ref.update_size();
273
274        Some(node_ref)
275    }
276
277    /// Deletes node with smallest key from the tree
278    ///
279    /// # Examples
280    /// ```
281    /// use rudac::tree::RedBlack;
282    ///
283    /// let mut rb_tree = RedBlack::<usize,usize>::init();
284    ///
285    /// rb_tree.insert(1,10);
286    /// rb_tree.insert(2,20);
287    /// rb_tree.insert(3,30);
288    /// rb_tree.insert(4,40);
289    ///
290    /// rb_tree.delete_min();
291    /// assert_eq!(rb_tree.get(&1), None);
292    ///
293    ///
294    /// rb_tree.delete_min();
295    /// assert_eq!(rb_tree.get(&2), None);
296    /// ```
297    pub fn delete_min(&mut self) {
298        if self.root.is_none() {
299            return;
300        }
301
302        let mut root_ref = self.root.take().unwrap();
303
304        if !Node::is_red(&root_ref.left_child) && !Node::is_red(&root_ref.right_child) {
305            root_ref.color = RED;
306        }
307
308        let mut root = RedBlack::_delete_min(Some(root_ref));
309
310        if !root.is_none() {
311            root_ref = root.unwrap();
312            root_ref.color = BLACK;
313            root = Some(root_ref)
314        }
315
316        self.root = root;
317    }
318
319    fn _delete_min(node: Option<Box<Node<K, V>>>) -> Option<Box<Node<K, V>>> {
320        if node.as_ref().unwrap().left_child.is_none() {
321            return None;
322        }
323
324        let mut node_ref = node.unwrap();
325
326        if !Node::is_red(&node_ref.left_child) && !Node::is_red(&node_ref.left_child().left_child) {
327            node_ref = RedBlack::move_red_left(node_ref);
328        }
329
330        node_ref.left_child = RedBlack::_delete_min(node_ref.left_child);
331
332        Some(RedBlack::balance(node_ref))
333    }
334
335    /// Deletes node with largest key from the tree
336    ///
337    /// # Examples
338    /// ```
339    /// use rudac::tree::RedBlack;
340    ///
341    /// let mut rb_tree = RedBlack::<usize,usize>::init();
342    ///
343    /// rb_tree.insert(1,10);
344    /// rb_tree.insert(2,20);
345    /// rb_tree.insert(3,30);
346    /// rb_tree.insert(4,40);
347    ///
348    /// rb_tree.delete_max();
349    /// assert_eq!(rb_tree.get(&4), None);
350    ///
351    ///
352    /// rb_tree.delete_max();
353    /// assert_eq!(rb_tree.get(&3), None);
354    /// ```
355    pub fn delete_max(&mut self) {
356        if self.root.is_none() {
357            return;
358        }
359
360        let mut root_ref = self.root.take().unwrap();
361
362        if !Node::is_red(&root_ref.left_child) && !Node::is_red(&root_ref.right_child) {
363            root_ref.color = RED;
364        }
365
366        let mut root = RedBlack::_delete_max(Some(root_ref));
367
368        if !root.is_none() {
369            root_ref = root.unwrap();
370            root_ref.color = BLACK;
371            root = Some(root_ref)
372        }
373
374        self.root = root;
375    }
376
377    fn _delete_max(node: Option<Box<Node<K, V>>>) -> Option<Box<Node<K, V>>> {
378        if node.is_none() {
379            return None;
380        }
381
382        let mut node_ref = node.unwrap();
383
384        if Node::is_red(&node_ref.left_child) {
385            node_ref = RedBlack::rotate_right(node_ref);
386        }
387
388        if node_ref.right_child.is_none() {
389            return None;
390        }
391
392        if !Node::is_red(&node_ref.right_child) && !Node::is_red(&node_ref.right_child().left_child)
393        {
394            node_ref = RedBlack::move_red_right(node_ref);
395        }
396
397        node_ref.right_child = RedBlack::_delete_max(node_ref.right_child);
398
399        Some(RedBlack::balance(node_ref))
400    }
401
402    /// Deletes the node containing the specified `key`
403    ///
404    /// # Arguments
405    /// * `key`: key of the node to be deleted from the tree
406    ///
407    /// # Examples
408    /// ```
409    /// use rudac::tree::RedBlack;
410    ///
411    /// let mut rb_tree = RedBlack::<usize,usize>::init();
412    ///
413    /// rb_tree.insert(1,10);
414    /// rb_tree.insert(2,20);
415    /// rb_tree.insert(3,30);
416    /// rb_tree.insert(4,40);
417    ///
418    /// rb_tree.delete(&1);
419    /// assert_eq!(rb_tree.get(&1), None);
420    /// ```
421    pub fn delete(&mut self, key: &K) {
422        if self.root.is_none() || !self.contains(key) {
423            return;
424        }
425
426        let mut root_ref = self.root.take().unwrap();
427
428        if !Node::is_red(&root_ref.left_child) && !Node::is_red(&root_ref.right_child) {
429            root_ref.color = RED;
430        }
431
432        let mut root = RedBlack::_delete(Some(root_ref), key);
433
434        if !root.is_none() {
435            root_ref = root.unwrap();
436            root_ref.color = BLACK;
437            root = Some(root_ref)
438        }
439
440        self.root = root;
441    }
442
443    fn _delete(node: Option<Box<Node<K, V>>>, key: &K) -> Option<Box<Node<K, V>>> {
444        if node.is_none() {
445            return None;
446        }
447
448        let mut node_ref = node.unwrap();
449
450        if *key < *node_ref.key() {
451            if !Node::is_red(&node_ref.left_child)
452                && Node::is_red(&node_ref.left_child().left_child)
453            {
454                node_ref = RedBlack::move_red_left(node_ref);
455            }
456            node_ref.left_child = RedBlack::_delete(node_ref.left_child, key);
457        } else {
458            if Node::is_red(&node_ref.left_child) {
459                node_ref = RedBlack::rotate_right(node_ref);
460            }
461            if *key == *node_ref.key() && node_ref.right_child.is_none() {
462                return None;
463            }
464            if !Node::is_red(&node_ref.right_child)
465                && !Node::is_red(&node_ref.right_child().left_child)
466            {
467                node_ref = RedBlack::move_red_right(node_ref);
468            }
469            if *key == *node_ref.key() {
470                let mut x = RedBlack::_min(&mut node_ref.right_child);
471                // swap keys
472                std::mem::swap(x.key_mut(), node_ref.key_mut());
473
474                // swap values
475                std::mem::swap(x.value_mut(), node_ref.value_mut());
476
477                node_ref.right_child = RedBlack::_delete_min(node_ref.right_child);
478            } else {
479                node_ref.right_child = RedBlack::_delete(node_ref.right_child, key);
480            }
481        }
482
483        Some(RedBlack::balance(node_ref))
484    }
485
486    fn _min(node: &mut Option<Box<Node<K, V>>>) -> Box<Node<K, V>> {
487        match node {
488            None => panic!("Called min on None node"),
489            Some(_node) => {
490                if _node.left_child.is_none() {
491                    return Box::new(Node::init(_node.get_key(), _node.get_value(), RED, 1));
492                } else {
493                    return RedBlack::_min(&mut _node.left_child);
494                }
495            }
496        }
497    }
498
499    /// Returns the height of the tree.
500    /// An empty tree has height -1 and a tree with one node has height 0
501    pub fn height(&self) -> i64 {
502        RedBlack::_height(&self.root)
503    }
504
505    fn _height(node: &Option<Box<Node<K, V>>>) -> i64 {
506        if node.is_none() {
507            return -1;
508        }
509
510        let node_ref = node.as_ref().unwrap();
511
512        return 1 + std::cmp::max(
513            RedBlack::_height(&node_ref.left_child),
514            RedBlack::_height(&node_ref.right_child),
515        );
516    }
517
518    /// Returns the largest key in the tree less than or equal to `key`
519    ///
520    /// # Arguments
521    /// * `key`: key to be searched for
522    ///
523    /// # Examples:
524    /// ```
525    /// use rudac::tree::RedBlack;
526    ///
527    /// let mut rb_tree = RedBlack::<usize,usize>::init();
528    ///
529    /// rb_tree.insert(1,10);
530    /// rb_tree.insert(3,20);
531    /// rb_tree.insert(5,30);
532    /// rb_tree.insert(7,40);
533    ///
534    /// assert_eq!(*rb_tree.floor(&2).unwrap(), 1);
535    /// assert_eq!(rb_tree.floor(&0), None);
536    /// ```
537    pub fn floor(&self, key: &K) -> Option<&K> {
538        RedBlack::_floor(&self.root, key)
539    }
540
541    fn _floor<'a>(node: &'a Option<Box<Node<K, V>>>, key: &K) -> Option<&'a K> {
542        if node.is_none() {
543            return None;
544        }
545
546        let node_ref = node.as_ref().unwrap();
547
548        if *key == *node_ref.key() {
549            return Some(node_ref.key());
550        }
551        if *key < *node_ref.key() {
552            return RedBlack::_floor(&node_ref.left_child, key);
553        }
554
555        let found_key = RedBlack::_floor(&node_ref.right_child, key);
556
557        if found_key.is_some() {
558            return found_key;
559        } else {
560            return Some(node_ref.key());
561        }
562    }
563
564    /// Returns the smallest key in the tree greater than or equal to `key`
565    ///
566    /// # Arguments
567    /// * `key`: key to be searched for
568    ///
569    /// # Examples:
570    /// ```
571    /// use rudac::tree::RedBlack;
572    ///
573    /// let mut rb_tree = RedBlack::<usize,usize>::init();
574    ///
575    /// rb_tree.insert(1,10);
576    /// rb_tree.insert(3,20);
577    /// rb_tree.insert(5,30);
578    /// rb_tree.insert(7,40);
579    ///
580    /// assert_eq!(*rb_tree.ceiling(&6).unwrap(), 7);
581    /// assert_eq!(rb_tree.ceiling(&8), None);
582    /// ```
583    pub fn ceiling(&self, key: &K) -> Option<&K> {
584        RedBlack::_ceiling(&self.root, key)
585    }
586
587    fn _ceiling<'a>(node: &'a Option<Box<Node<K, V>>>, key: &K) -> Option<&'a K> {
588        if node.is_none() {
589            return None;
590        }
591
592        let node_ref = node.as_ref().unwrap();
593
594        if *key == *node_ref.key() {
595            return Some(node_ref.key());
596        }
597        if *key > *node_ref.key() {
598            return RedBlack::_ceiling(&node_ref.right_child, key);
599        }
600
601        let found_key = RedBlack::_ceiling(&node_ref.left_child, key);
602
603        if found_key.is_some() {
604            return found_key;
605        } else {
606            return Some(node_ref.key());
607        }
608    }
609
610    /// Returns the kth smallest key and its associated value in the tree
611    ///
612    /// # Arguments
613    /// * `k`: the order statistic
614    ///
615    /// # Panics
616    /// * panics if k is not in range: 0 <= k <= size - 1
617    ///
618    /// # Examples
619    /// ```
620    /// use rudac::tree::RedBlack;
621    ///
622    /// let mut rb_tree = RedBlack::<usize,usize>::init();
623    ///
624    /// rb_tree.insert(1,10);
625    /// rb_tree.insert(3,20);
626    /// rb_tree.insert(5,30);
627    /// rb_tree.insert(7,40);
628    ///
629    /// let (key, value) = rb_tree.select(1).unwrap();
630    /// assert_eq!(*key, 3);
631    /// assert_eq!(*value, 20);
632    /// ```
633    pub fn select(&self, k: usize) -> Option<(&K, &V)> {
634        if k > self.size() {
635            panic!("K must be in range 0 <= k <= size - 1");
636        }
637        RedBlack::_select(&self.root, k)
638    }
639
640    fn _select(node: &Option<Box<Node<K, V>>>, k: usize) -> Option<(&K, &V)> {
641        if node.is_none() {
642            return None;
643        }
644
645        let node_ref = node.as_ref().unwrap();
646
647        let left_size = Node::size(&node_ref.left_child);
648
649        if left_size > k {
650            return RedBlack::_select(&node_ref.left_child, k);
651        } else if left_size < k {
652            return RedBlack::_select(&node_ref.right_child, k - left_size - 1);
653        } else {
654            return Some((node_ref.key(), node_ref.value()));
655        }
656    }
657
658    /// Returns the smallest key and its associated value in the tree
659    ///
660    /// # Examples
661    /// ```
662    /// use rudac::tree::RedBlack;
663    ///
664    /// let mut rb_tree = RedBlack::<usize,usize>::init();
665    ///
666    /// rb_tree.insert(1,10);
667    /// rb_tree.insert(3,20);
668    /// rb_tree.insert(5,30);
669    /// rb_tree.insert(7,40);
670    ///
671    /// let (key, value) = rb_tree.min().unwrap();
672    /// assert_eq!(*key, 1);
673    /// assert_eq!(*value, 10);
674    ///
675    /// rb_tree.delete_min();
676    ///
677    /// let (key, value) = rb_tree.min().unwrap();
678    /// assert_eq!(*key, 3);
679    /// assert_eq!(*value, 20);
680    /// ```
681    pub fn min(&self) -> Option<(&K, &V)> {
682        self.select(0)
683    }
684
685    /// Returns the largest key and its associated value in the tree
686    ///
687    /// # Examples
688    /// ```
689    /// use rudac::tree::RedBlack;
690    ///
691    /// let mut rb_tree = RedBlack::<usize,usize>::init();
692    ///
693    /// rb_tree.insert(1,10);
694    /// rb_tree.insert(3,20);
695    /// rb_tree.insert(5,30);
696    /// rb_tree.insert(7,40);
697    ///
698    /// let (key, value) = rb_tree.max().unwrap();
699    /// assert_eq!(*key, 7);
700    /// assert_eq!(*value, 40);
701    ///
702    /// rb_tree.delete_max();
703    ///
704    /// let (key, value) = rb_tree.max().unwrap();
705    /// assert_eq!(*key, 5);
706    /// assert_eq!(*value, 30);
707    /// ```
708    pub fn max(&self) -> Option<(&K, &V)> {
709        self.select(self.size() - 1)
710    }
711
712    /// Returns the number of keys in the symbol table strictly less than `key`
713    ///
714    /// # Arguments
715    /// * `key`: key to be searched for
716    ///
717    /// # Examples
718    /// ```
719    /// use rudac::tree::RedBlack;
720    ///
721    /// let mut rb_tree = RedBlack::<usize, usize>::init();
722    ///
723    /// for i in 1..100 {
724    ///     rb_tree.insert(i, i);
725    /// }
726    ///
727    /// assert_eq!(rb_tree.rank(&99), 98);
728    /// ```
729    pub fn rank(&self, key: &K) -> usize {
730        RedBlack::_rank(&self.root, key)
731    }
732
733    fn _rank(node: &Option<Box<Node<K, V>>>, key: &K) -> usize {
734        if node.is_none() {
735            return 0;
736        }
737
738        let node_ref = node.as_ref().unwrap();
739
740        if *key < *node_ref.key() {
741            return RedBlack::_rank(&node_ref.left_child, key);
742        } else if *key > *node_ref.key() {
743            return 1
744                + Node::size(&node_ref.left_child)
745                + RedBlack::_rank(&node_ref.right_child, key);
746        } else {
747            return Node::size(&node_ref.left_child);
748        }
749    }
750
751    /// Returns all keys in the tree following an in-order traversal.
752    /// Therefore keys are sorted from smallest to largest
753    ///
754    /// # Examples
755    /// ```
756    /// use rudac::tree::RedBlack;
757    ///
758    /// let mut rb_tree = RedBlack::<usize, usize>::init();
759    ///
760    /// for i in (1..100).rev() {
761    ///     rb_tree.insert(i, i);
762    /// }
763    ///
764    /// let mut i = 1;
765    /// // keys are sorted: [1, 2, 3,..., 99]
766    /// for key in rb_tree.keys() {
767    ///     assert!(*key == i);
768    ///     i += 1;
769    /// }
770    /// ```
771    pub fn keys(&self) -> Vec<&K> {
772        let mut keys: Vec<&K> = Vec::new();
773
774        RedBlack::_keys_in_order(&self.root, &mut keys);
775
776        keys
777    }
778
779    fn _keys_in_order<'a>(node: &'a Option<Box<Node<K, V>>>, keys: &mut Vec<&'a K>) {
780        if node.is_none() {
781            return;
782        }
783
784        let node_ref = node.as_ref().unwrap();
785        RedBlack::_keys_in_order(&node_ref.left_child, keys);
786        keys.push(node_ref.key());
787        RedBlack::_keys_in_order(&node_ref.right_child, keys);
788    }
789
790    /// Returns all keys in the tree following a level-order traversal
791    pub fn keys_in_level_order(&self) -> Vec<&K> {
792        let mut keys: Vec<&K> = Vec::new();
793
794        RedBlack::_keys_in_level_order(&self.root, &mut keys);
795
796        keys
797    }
798
799    fn _keys_in_level_order<'a>(node: &'a Option<Box<Node<K, V>>>, keys: &mut Vec<&'a K>) {
800        if node.is_none() {
801            return;
802        }
803
804        let mut queue = VecDeque::<&Option<Box<Node<K, V>>>>::with_capacity(Node::size(node));
805        queue.push_back(node);
806
807        while !queue.is_empty() {
808            let current_node = queue.pop_front().unwrap().as_ref().unwrap();
809            keys.push(current_node.key());
810
811            if !current_node.left_child.is_none() {
812                queue.push_back(&current_node.left_child);
813            }
814            if !current_node.right_child.is_none() {
815                queue.push_back(&current_node.right_child);
816            }
817        }
818    }
819
820    /// Returns all keys in the symbol table between `low_key`(inclusive) and `high_key`(exclusive)
821    ///
822    /// # Arguments
823    /// * `low_key`: lowest key of the range
824    /// * `high_key`: highest key of the range
825    ///
826    /// # Examples
827    /// ```
828    /// use rudac::tree::RedBlack;
829    ///
830    /// let mut rb_tree = RedBlack::<usize, usize>::init();
831    ///
832    /// for i in (1..100).rev() {
833    ///     rb_tree.insert(i, i);
834    /// }
835    ///
836    /// let keys = rb_tree.keys_between(&1, &99);
837    ///
838    /// assert_eq!(keys.len(), 98);
839    /// ```
840    pub fn keys_between(&self, low_key: &K, high_key: &K) -> Vec<&K> {
841        let mut keys: Vec<&K> = Vec::new();
842
843        RedBlack::_keys_between(&self.root, low_key, high_key, &mut keys);
844
845        keys
846    }
847
848    fn _keys_between<'a>(
849        node: &'a Option<Box<Node<K, V>>>,
850        low_key: &K,
851        high_key: &K,
852        keys: &mut Vec<&'a K>,
853    ) {
854        if node.is_none() {
855            return;
856        }
857
858        let node_ref = node.as_ref().unwrap();
859        if *low_key < *node_ref.key() {
860            RedBlack::_keys_between(&node_ref.left_child, low_key, high_key, keys);
861        }
862        if *low_key <= *node_ref.key() && *node_ref.key() < *high_key {
863            keys.push(node_ref.key());
864        }
865        if *high_key > *node_ref.key() {
866            RedBlack::_keys_between(&node_ref.right_child, low_key, high_key, keys);
867        }
868    }
869
870    /// Returns the number of keys in the tree between `low_key`(inclusive) and `high_key`(exclusive)
871    ///
872    /// # Arguments
873    /// * `low_key`: lowest key of the range
874    /// * `high_key`: highest key of the range
875    ///
876    /// # Examples
877    /// ```
878    /// use rudac::tree::RedBlack;
879    ///
880    /// let mut rb_tree = RedBlack::<usize, usize>::init();
881    ///
882    /// for i in (1..100).rev() {
883    ///     rb_tree.insert(i, i);
884    /// }
885    ///
886    /// let keys = rb_tree.size_between(&1, &99);
887    ///
888    /// assert_eq!(keys, 98);
889    /// ```
890    pub fn size_between(&self, low_key: &K, high_key: &K) -> usize {
891        if self.is_empty() {
892            return 0;
893        }
894        if *low_key > *high_key {
895            return 0;
896        }
897
898        return self.rank(high_key) - self.rank(low_key);
899    }
900
901    fn rotate_left(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
902        let mut y = node.right_child.unwrap();
903        node.right_child = y.left_child;
904
905        // update colors
906        y.color = node.color;
907        node.color = RED;
908
909        // update y size
910        y.size = node.size;
911
912        // update node size
913        node.update_size();
914
915        // add node as left child of y
916        y.left_child = Some(node);
917        y
918    }
919
920    fn rotate_right(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
921        let mut y = node.left_child.unwrap();
922        node.left_child = y.right_child;
923
924        // update colors
925        y.color = node.color;
926        node.color = RED;
927
928        // update y size
929        y.size = node.size;
930
931        // update node size
932        node.update_size();
933
934        // add node as right child of y
935        y.right_child = Some(node);
936        y
937    }
938
939    fn flip_colors(node: &mut Box<Node<K, V>>) {
940        node.color = !node.color;
941        // flip left child color
942        let mut left_child = node.left_child.take().unwrap();
943        left_child.color = !left_child.color;
944
945        // flip right child color
946        let mut right_child = node.right_child.take().unwrap();
947        right_child.color = !right_child.color;
948
949        node.left_child = Some(left_child);
950        node.right_child = Some(right_child);
951    }
952
953    fn move_red_left(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
954        RedBlack::flip_colors(&mut node);
955
956        if Node::is_red(&node.right_child().left_child) {
957            let mut right_child = node.right_child.take().unwrap();
958            right_child = RedBlack::rotate_right(right_child);
959
960            node.right_child = Some(right_child);
961
962            node = RedBlack::rotate_left(node);
963            RedBlack::flip_colors(&mut node);
964        }
965
966        node
967    }
968
969    fn move_red_right(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
970        RedBlack::flip_colors(&mut node);
971
972        if Node::is_red(&node.left_child().left_child) {
973            node = RedBlack::rotate_right(node);
974            RedBlack::flip_colors(&mut node);
975        }
976
977        node
978    }
979
980    fn balance(mut node: Box<Node<K, V>>) -> Box<Node<K, V>> {
981        if Node::is_red(&node.right_child) {
982            node = RedBlack::rotate_left(node);
983        }
984        if Node::is_red(&node.left_child) && Node::is_red(&node.left_child().left_child) {
985            node = RedBlack::rotate_right(node);
986        }
987        if Node::is_red(&node.left_child) && Node::is_red(&node.right_child) {
988            RedBlack::flip_colors(&mut node);
989        }
990
991        node.update_size();
992
993        node
994    }
995}
996
997#[cfg(test)]
998mod tests {
999    use super::*;
1000
1001    fn is_23<K: std::cmp::Ord, V>(node: &Option<Box<Node<K, V>>>, is_root: bool) -> bool {
1002        if node.is_none() {
1003            return true;
1004        }
1005        let node_ref = node.as_ref().unwrap();
1006
1007        if Node::is_red(&node_ref.right_child) {
1008            return false;
1009        }
1010        if !is_root && Node::is_red(node) && Node::is_red(&node_ref.left_child) {
1011            return false;
1012        }
1013
1014        return is_23(&node_ref.left_child, false) && is_23(&node_ref.right_child, false);
1015    }
1016
1017    fn is_bst<K: std::cmp::Ord, V>(
1018        node: &Option<Box<Node<K, V>>>,
1019        min: Option<&K>,
1020        max: Option<&K>,
1021    ) -> bool {
1022        if node.is_none() {
1023            return true;
1024        }
1025
1026        let node_ref = node.as_ref().unwrap();
1027        if !min.is_none() && *node_ref.key() <= **min.as_ref().unwrap() {
1028            return false;
1029        }
1030        if !max.is_none() && *node_ref.key() >= **max.as_ref().unwrap() {
1031            return false;
1032        }
1033
1034        return is_bst(&node_ref.left_child, min, Some(node_ref.key()))
1035            && is_bst(&node_ref.right_child, Some(node_ref.key()), max);
1036    }
1037
1038    fn is_size_consistent<K: std::cmp::Ord, V>(node: &Option<Box<Node<K, V>>>) -> bool {
1039        if node.is_none() {
1040            return true;
1041        }
1042        let node_ref = node.as_ref().unwrap();
1043
1044        if Node::size(node)
1045            != Node::size(&node_ref.left_child) + Node::size(&node_ref.right_child) + 1
1046        {
1047            return false;
1048        }
1049
1050        return is_size_consistent(&node_ref.left_child)
1051            && is_size_consistent(&node_ref.right_child);
1052    }
1053
1054    fn is_rank_consistent<K: std::cmp::Ord, V>(rb_tree: &RedBlack<K, V>) -> bool {
1055        for i in 0..Node::size(&rb_tree.root) {
1056            if i != rb_tree.rank(rb_tree.select(i).unwrap().0) {
1057                return false;
1058            }
1059        }
1060
1061        for key in rb_tree.keys() {
1062            if *rb_tree.select(rb_tree.rank(key)).unwrap().0 != *key {
1063                return false;
1064            }
1065        }
1066
1067        true
1068    }
1069
1070    #[test]
1071    fn tree_rb_init() {
1072        let rb_tree = RedBlack::<usize, usize>::init();
1073
1074        assert!(is_23(&rb_tree.root, true));
1075        assert!(is_bst(&rb_tree.root, None, None));
1076        assert!(is_size_consistent(&rb_tree.root));
1077        assert!(is_rank_consistent(&rb_tree));
1078    }
1079
1080    #[test]
1081    fn tree_rb_insert_1() {
1082        let mut rb_tree = RedBlack::<usize, usize>::init();
1083
1084        rb_tree.insert(1, 1);
1085
1086        assert!(is_23(&rb_tree.root, true));
1087        assert!(is_bst(&rb_tree.root, None, None));
1088        assert!(is_size_consistent(&rb_tree.root));
1089        assert!(is_rank_consistent(&rb_tree));
1090    }
1091
1092    #[test]
1093    fn tree_rb_insert_2() {
1094        let mut rb_tree = RedBlack::<usize, usize>::init();
1095
1096        rb_tree.insert(4, 1);
1097        rb_tree.insert(3, 1);
1098        rb_tree.insert(2, 1);
1099        rb_tree.insert(1, 1);
1100
1101        assert!(is_23(&rb_tree.root, true));
1102        assert!(is_bst(&rb_tree.root, None, None));
1103        assert!(is_size_consistent(&rb_tree.root));
1104        assert!(is_rank_consistent(&rb_tree));
1105    }
1106
1107    #[test]
1108    fn tree_rb_insert_3() {
1109        let mut rb_tree = RedBlack::<usize, usize>::init();
1110
1111        for i in (0..100).rev() {
1112            rb_tree.insert(i, i);
1113        }
1114
1115        assert!(is_23(&rb_tree.root, true));
1116        assert!(is_bst(&rb_tree.root, None, None));
1117        assert!(is_size_consistent(&rb_tree.root));
1118        assert!(is_rank_consistent(&rb_tree));
1119    }
1120
1121    #[test]
1122    fn tree_rb_delete_1() {
1123        let mut rb_tree = RedBlack::<usize, usize>::init();
1124
1125        rb_tree.delete(&1);
1126
1127        assert!(is_23(&rb_tree.root, true));
1128        assert!(is_bst(&rb_tree.root, None, None));
1129        assert!(is_size_consistent(&rb_tree.root));
1130        assert!(is_rank_consistent(&rb_tree));
1131    }
1132
1133    #[test]
1134    fn tree_rb_delete_2() {
1135        let mut rb_tree = RedBlack::<usize, usize>::init();
1136
1137        rb_tree.insert(4, 1);
1138        rb_tree.insert(3, 1);
1139        rb_tree.insert(2, 1);
1140        rb_tree.insert(1, 1);
1141
1142        rb_tree.delete(&1);
1143        assert_eq!(rb_tree.get(&1), None);
1144
1145        rb_tree.delete(&3);
1146        assert_eq!(rb_tree.get(&3), None);
1147
1148        rb_tree.delete(&2);
1149        assert_eq!(rb_tree.get(&2), None);
1150
1151        rb_tree.delete(&4);
1152        assert_eq!(rb_tree.get(&4), None);
1153
1154        assert!(is_23(&rb_tree.root, true));
1155        assert!(is_bst(&rb_tree.root, None, None));
1156        assert!(is_size_consistent(&rb_tree.root));
1157        assert!(is_rank_consistent(&rb_tree));
1158    }
1159
1160    #[test]
1161    fn tree_rb_is_empty() {
1162        let mut rb_tree = RedBlack::<usize, usize>::init();
1163
1164        assert!(rb_tree.is_empty());
1165
1166        rb_tree.insert(1, 1);
1167
1168        assert!(!rb_tree.is_empty());
1169    }
1170
1171    #[test]
1172    fn tree_rb_size_1() {
1173        let mut rb_tree = RedBlack::<usize, usize>::init();
1174        assert_eq!(rb_tree.size(), 0);
1175
1176        rb_tree.insert(1, 1);
1177        assert_eq!(rb_tree.size(), 1);
1178
1179        rb_tree.insert(2, 1);
1180        assert_eq!(rb_tree.size(), 2);
1181
1182        rb_tree.insert(3, 1);
1183        assert_eq!(rb_tree.size(), 3);
1184
1185        rb_tree.insert(4, 1);
1186        assert_eq!(rb_tree.size(), 4);
1187    }
1188
1189    #[test]
1190    fn tree_rb_contains_1() {
1191        let mut rb_tree = RedBlack::<usize, usize>::init();
1192        assert!(!rb_tree.contains(&1));
1193
1194        rb_tree.insert(1, 2);
1195        assert!(rb_tree.contains(&1));
1196
1197        rb_tree.insert(1, 3);
1198        assert!(rb_tree.contains(&1));
1199    }
1200
1201    #[test]
1202    fn tree_rb_get_1() {
1203        let mut rb_tree = RedBlack::<usize, usize>::init();
1204        assert_eq!(rb_tree.get(&1), None);
1205
1206        rb_tree.insert(1, 2);
1207        assert_eq!(*rb_tree.get(&1).unwrap(), 2);
1208
1209        assert!(is_23(&rb_tree.root, true));
1210        assert!(is_bst(&rb_tree.root, None, None));
1211        assert!(is_size_consistent(&rb_tree.root));
1212        assert!(is_rank_consistent(&rb_tree));
1213    }
1214
1215    #[test]
1216    fn tree_rb_get_2() {
1217        let mut rb_tree = RedBlack::<usize, usize>::init();
1218
1219        for i in (0..100).rev() {
1220            rb_tree.insert(i, i);
1221        }
1222
1223        for i in (0..100).rev() {
1224            assert_eq!(*rb_tree.get(&i).unwrap(), i);
1225        }
1226
1227        assert!(is_23(&rb_tree.root, true));
1228        assert!(is_bst(&rb_tree.root, None, None));
1229        assert!(is_size_consistent(&rb_tree.root));
1230        assert!(is_rank_consistent(&rb_tree));
1231    }
1232
1233    #[test]
1234    fn tree_rb_get_3() {
1235        let mut rb_tree = RedBlack::<usize, usize>::init();
1236
1237        for i in (0..100).rev() {
1238            rb_tree.insert(i, i);
1239        }
1240
1241        for i in (0..100).rev() {
1242            assert_eq!(*rb_tree.get(&i).unwrap(), i);
1243        }
1244
1245        for i in (0..100).rev() {
1246            rb_tree.insert(i, i + 1);
1247        }
1248
1249        for i in (0..100).rev() {
1250            assert_eq!(*rb_tree.get(&i).unwrap(), i + 1);
1251        }
1252
1253        assert!(is_23(&rb_tree.root, true));
1254        assert!(is_bst(&rb_tree.root, None, None));
1255        assert!(is_size_consistent(&rb_tree.root));
1256        assert!(is_rank_consistent(&rb_tree));
1257    }
1258
1259    #[test]
1260    fn tree_rb_delete_min_1() {
1261        let mut rb_tree = RedBlack::<usize, usize>::init();
1262
1263        rb_tree.delete_min();
1264
1265        assert!(is_23(&rb_tree.root, true));
1266        assert!(is_bst(&rb_tree.root, None, None));
1267        assert!(is_size_consistent(&rb_tree.root));
1268        assert!(is_rank_consistent(&rb_tree));
1269    }
1270
1271    #[test]
1272    fn tree_rb_delete_min_2() {
1273        let mut rb_tree = RedBlack::<usize, usize>::init();
1274
1275        for i in (0..100).rev() {
1276            rb_tree.insert(i, i);
1277        }
1278
1279        for i in 0..100 {
1280            rb_tree.delete_min();
1281            assert_eq!(rb_tree.get(&i), None);
1282        }
1283
1284        assert!(is_23(&rb_tree.root, true));
1285        assert!(is_bst(&rb_tree.root, None, None));
1286        assert!(is_size_consistent(&rb_tree.root));
1287        assert!(is_rank_consistent(&rb_tree));
1288    }
1289
1290    #[test]
1291    fn tree_rb_delete_max_1() {
1292        let mut rb_tree = RedBlack::<usize, usize>::init();
1293
1294        rb_tree.delete_max();
1295
1296        assert!(is_23(&rb_tree.root, true));
1297        assert!(is_bst(&rb_tree.root, None, None));
1298        assert!(is_size_consistent(&rb_tree.root));
1299        assert!(is_rank_consistent(&rb_tree));
1300    }
1301
1302    #[test]
1303    fn tree_rb_delete_max_2() {
1304        let mut rb_tree = RedBlack::<usize, usize>::init();
1305
1306        for i in (0..100).rev() {
1307            rb_tree.insert(i, i);
1308        }
1309
1310        for i in (0..100).rev() {
1311            rb_tree.delete_max();
1312            assert_eq!(rb_tree.get(&i), None);
1313        }
1314
1315        assert!(is_23(&rb_tree.root, true));
1316        assert!(is_bst(&rb_tree.root, None, None));
1317        assert!(is_size_consistent(&rb_tree.root));
1318        assert!(is_rank_consistent(&rb_tree));
1319    }
1320
1321    #[test]
1322    fn tree_rb_floor_1() {
1323        let mut rb_tree = RedBlack::<usize, usize>::init();
1324
1325        for i in (0..100).step_by(2) {
1326            rb_tree.insert(i, i);
1327        }
1328
1329        for i in (1..100).step_by(2) {
1330            assert_eq!(*rb_tree.floor(&i).unwrap(), i - 1);
1331        }
1332
1333        assert!(is_23(&rb_tree.root, true));
1334        assert!(is_bst(&rb_tree.root, None, None));
1335        assert!(is_size_consistent(&rb_tree.root));
1336        assert!(is_rank_consistent(&rb_tree));
1337    }
1338
1339    #[test]
1340    fn tree_rb_ceiling_1() {
1341        let mut rb_tree = RedBlack::<usize, usize>::init();
1342
1343        for i in (0..100).step_by(2) {
1344            rb_tree.insert(i, i);
1345        }
1346
1347        for i in (1..99).step_by(2) {
1348            assert_eq!(*rb_tree.ceiling(&i).unwrap(), i + 1);
1349        }
1350
1351        assert!(is_23(&rb_tree.root, true));
1352        assert!(is_bst(&rb_tree.root, None, None));
1353        assert!(is_size_consistent(&rb_tree.root));
1354        assert!(is_rank_consistent(&rb_tree));
1355    }
1356
1357    #[test]
1358    fn tree_rb_select_1() {
1359        let mut rb_tree = RedBlack::<usize, usize>::init();
1360
1361        for i in (0..100).rev() {
1362            rb_tree.insert(i, i);
1363        }
1364
1365        for i in 0..100 {
1366            let result = rb_tree.select(i).unwrap();
1367            assert_eq!((*result.0, *result.1), (i, i));
1368        }
1369
1370        assert!(is_23(&rb_tree.root, true));
1371        assert!(is_bst(&rb_tree.root, None, None));
1372        assert!(is_size_consistent(&rb_tree.root));
1373        assert!(is_rank_consistent(&rb_tree));
1374    }
1375
1376    #[test]
1377    fn tree_rb_rank_1() {
1378        let mut rb_tree = RedBlack::<usize, usize>::init();
1379
1380        for i in (1..100).rev() {
1381            rb_tree.insert(i, i);
1382        }
1383
1384        for i in 1..100 {
1385            assert_eq!(rb_tree.rank(&i), i - 1);
1386        }
1387
1388        assert!(is_23(&rb_tree.root, true));
1389        assert!(is_bst(&rb_tree.root, None, None));
1390        assert!(is_size_consistent(&rb_tree.root));
1391        assert!(is_rank_consistent(&rb_tree));
1392    }
1393
1394    #[test]
1395    fn tree_rb_keys_1() {
1396        let mut rb_tree = RedBlack::<usize, usize>::init();
1397
1398        for i in (1..100).rev() {
1399            rb_tree.insert(i, i);
1400        }
1401
1402        let mut i = 1;
1403        for key in rb_tree.keys() {
1404            assert!(*key == i);
1405            i += 1;
1406        }
1407
1408        assert!(is_23(&rb_tree.root, true));
1409        assert!(is_bst(&rb_tree.root, None, None));
1410        assert!(is_size_consistent(&rb_tree.root));
1411        assert!(is_rank_consistent(&rb_tree));
1412    }
1413
1414    #[test]
1415    fn tree_rb_keys_between_1() {
1416        let mut rb_tree = RedBlack::<usize, usize>::init();
1417
1418        for i in (1..100).rev() {
1419            rb_tree.insert(i, i);
1420        }
1421
1422        for i in 1..100 {
1423            assert_eq!(rb_tree.keys_between(&i, &99).len(), 99 - i);
1424        }
1425
1426        assert!(is_23(&rb_tree.root, true));
1427        assert!(is_bst(&rb_tree.root, None, None));
1428        assert!(is_size_consistent(&rb_tree.root));
1429        assert!(is_rank_consistent(&rb_tree));
1430    }
1431
1432    #[test]
1433    fn tree_rb_size_between_1() {
1434        let mut rb_tree = RedBlack::<usize, usize>::init();
1435
1436        for i in (1..100).rev() {
1437            rb_tree.insert(i, i);
1438        }
1439
1440        for i in 1..100 {
1441            assert_eq!(rb_tree.size_between(&i, &100), 100 - i);
1442        }
1443
1444        assert!(is_23(&rb_tree.root, true));
1445        assert!(is_bst(&rb_tree.root, None, None));
1446        assert!(is_size_consistent(&rb_tree.root));
1447        assert!(is_rank_consistent(&rb_tree));
1448    }
1449}
rudac/tree/rb.rs

rudac/tree/
rb.rs
