rust_black_trees/
node.rs

1use std::cmp::max;
2use std::ops::Not;
3
4#[derive(Debug, Clone, Copy)]
5pub enum Color {
6    Red,
7    Black,
8}
9
10#[derive(Debug, Clone, Copy)]
11pub enum Side {
12    Left,
13    Right,
14}
15
16impl Not for Side {
17    type Output = Side;
18    fn not(self) -> Self::Output {
19        match self {
20            Side::Left => Side::Right,
21            Side::Right => Side::Left,
22        }
23    }
24}
25
26#[derive(Debug)]
27pub struct DepthNode<T> {
28    pub value: T,
29    pub ptr: usize,
30    pub parent: Option<usize>,
31    pub lchild: Option<usize>,
32    pub rchild: Option<usize>,
33    pub height: usize,
34}
35
36pub fn paint(fg: usize, bg: usize) -> String {
37    let esc = char::from(0x1b);
38    format!("{}[{};{}m", esc, fg, bg)
39}
40pub fn endpaint() -> String {
41    let esc = char::from(0x1b);
42    format!("{}[0m", esc)
43}
44
45
46pub trait Node<T> {
47    // Base methods
48    fn get_value(&self) -> &T;
49    fn get(&self, i: usize) -> &Self;
50    fn get_mut(&self, i: usize) -> &mut Self;
51    fn location(&self) -> usize;
52    fn get_parent(&self) -> Option<usize>;
53    fn set_parent(&mut self, p: Option<usize>);
54    fn get_child(&self, side: Side) -> Option<usize>;
55    fn set_child(&mut self, child: usize, side: Side);
56    fn set_child_opt(&mut self, child: Option<usize>, side: Side);
57    fn to_self_string(&self) -> String;
58    fn to_self_string_display(&self) -> (String, usize);
59    fn is(&self, val: &T) -> bool;
60    fn greater(&self, val: &T) -> bool;
61    fn lesser(&self, val: &T) -> bool;
62
63    fn to_string(&self) -> String {
64        let mut m_str = format!("({}", self.to_self_string());
65        m_str = m_str
66            + " "
67            + &(if let Some(child) = self.get_child(Side::Left) {
68                self.get(child).to_string()
69            } else {
70                String::from("()")
71            });
72        m_str = m_str
73            + " "
74            + &(if let Some(child) = self.get_child(Side::Right) {
75                self.get(child).to_string()
76            } else {
77                String::from("()")
78            });
79        m_str + ")"
80    }
81
82    fn to_pretty_string(&self, indent: usize) -> String {
83        let i = indent * 2;
84        let mut m_str = format!("({}", self.to_self_string());
85        m_str = m_str
86            + &(if let Some(child) = self.get_child(Side::Left) {
87                "\n".to_owned() + &" ".repeat(i) + &self.get(child).to_pretty_string(indent + 1)
88            } else {
89                String::from(" ()")
90            });
91        m_str = m_str
92            + &(if let Some(child) = self.get_child(Side::Right) {
93                "\n".to_owned() + &" ".repeat(i) + &self.get(child).to_pretty_string(indent + 1)
94            } else {
95                String::from(" ()")
96            });
97        m_str + ")"
98    }
99
100    fn get_height(&self) -> usize {
101        let f = |c| Some(1 + self.get(c).get_height());
102        max(
103            self.get_child(Side::Left).and_then(f).unwrap_or(1),
104            self.get_child(Side::Right).and_then(f).unwrap_or(1),
105        )
106    }
107
108    fn get_depth(&self) -> usize {
109        let f = |c| Some(1 + self.get(c).get_depth());
110        self.get_parent().and_then(f).unwrap_or(0)
111    }
112
113    fn get_size(&self) -> usize {
114        let f = |c| Some(self.get(c).get_size());
115
116        1 + self.get_child(Side::Left).and_then(f).unwrap_or(0)
117            + self.get_child(Side::Right).and_then(f).unwrap_or(0)
118    }
119
120    fn get_leaf_count(&self) -> usize {
121        let f = |c| Some(self.get(c).get_leaf_count());
122        let val = self.get_child(Side::Left).and_then(f).unwrap_or(0) + self.get_child(Side::Right).and_then(f).unwrap_or(0);
123        if val == 0 {
124            1
125        } else {
126            val
127        }
128    }
129
130    fn find_min(&self) -> usize {
131        if let Some(l) = self.get_child(Side::Left) {
132            self.get(l).find_min()
133        } else {
134            self.location()
135        }
136    }
137
138    fn side(&self) -> Side {
139        if self.is_child(Side::Left) {
140            Side::Left
141        } else {
142            Side::Right
143        }
144    }
145
146    fn get_sibling(&self) -> Option<usize> {
147        if let Some(p) = self.get_parent() {
148            let parent = self.get(p);
149            if self.is_child(Side::Left) {
150                parent.get_child(Side::Right)
151            } else if self.is_child(Side::Right) {
152                parent.get_child(Side::Left)
153            } else {
154                None
155            }
156        } else {
157            None
158        }
159    }
160
161    fn get_uncle(&self) -> Option<usize> {
162        self.get_parent()
163            .and_then(|p| Some(self.get(p)))
164            .and_then(|p| p.get_sibling())
165    }
166
167    fn is_child(&self, side: Side) -> bool {
168        if let Some(p) = self.get_parent() {
169            let parent = self.get(p);
170            parent.get_child(side).is_some() && parent.get_child(side).unwrap() == self.location()
171        } else {
172            false
173        }
174    }
175}
176
177#[cfg(test)]
178mod tests {
179    use super::*;
180    use super::Node;
181    use crate::rbtree::ColorNode;
182    use crate::rbtree::ColoredNode;
183    use std::rc::Rc;
184    use std::cell::RefCell;
185
186    fn attach_child(data: &mut Vec<ColorNode<i32>>, p: usize, c: usize, side: Side) {
187        let par = &mut data[p];
188        par.set_child(c, side)
189    }
190    /*
191    ========
192    = Data =
193    ========
194    [0, 1, 2, 3, 4, 5] // ptrs
195    [5, 4, 6, 0, 1, 7] // vals
196    ([P:None C:Black V:5]
197         ([P:Some(0) C:Black V:1]
198              ([P:Some(4) C:Black V:0])
199              ([P:Some(4) C:Black V:4]))
200         ([P:Some(0) C:Black V:6]
201              ([P:Some(2) C:Black V:7])))
202    */
203    fn make_fake_tree_node() -> Rc<RefCell<Vec<ColorNode<i32>>>> {
204        let rc = Rc::new(RefCell::new(Vec::new()));
205        let v = vec![
206            ColorNode::new(5, 0, rc.clone()),
207            ColorNode::new(4, 1, rc.clone()),
208            ColorNode::new(6, 2, rc.clone()),
209            ColorNode::new(0, 3, rc.clone()),
210            ColorNode::new(1, 4, rc.clone()),
211            ColorNode::new(7, 5, rc.clone()),
212        ];
213        (*rc).replace(v);
214        let ptrs: Vec<usize> = (*rc).borrow_mut().iter().map(|v| v.ptr).collect();
215        attach_child(&mut (*rc).borrow_mut(), ptrs[0], ptrs[4], Side::Left);
216        attach_child(&mut (*rc).borrow_mut(), ptrs[0], ptrs[2], Side::Right);
217        attach_child(&mut (*rc).borrow_mut(), ptrs[4], ptrs[3], Side::Left);
218        attach_child(&mut (*rc).borrow_mut(), ptrs[4], ptrs[1], Side::Right);
219        attach_child(&mut (*rc).borrow_mut(), ptrs[2], ptrs[5], Side::Right);
220        rc
221    }
222
223    #[test]
224    fn print_tree_test() {
225        let rc = make_fake_tree_node();
226        let data = rc.borrow_mut();
227        let root = &data[0];
228        assert_eq!(root.to_string(), "([P:None C:Black V:5] ([P:Some(0) C:Black V:1] ([P:Some(4) C:Black V:0] () ()) ([P:Some(4) C:Black V:4] () ())) ([P:Some(0) C:Black V:6] () ([P:Some(2) C:Black V:7] () ())))");
229    }
230
231    #[test]
232    fn get_child_test() {
233        let rc = make_fake_tree_node();
234        let data = rc.borrow_mut();
235        let root = &data[0];
236
237        assert_eq!(root.get_child(Side::Left), Some(4));
238        assert_eq!(root.get_child(Side::Right), Some(2));
239        assert_eq!(data[2].get_child(Side::Right), Some(5));
240        assert_eq!(
241            data[root.get_child(Side::Left).unwrap()].get_child(Side::Right),
242            Some(1)
243        );
244        assert_eq!(
245            data[root.get_child(Side::Left).unwrap()].get_child(Side::Left),
246            Some(3)
247        );
248
249        assert_eq!(root.value, 5);
250        assert_eq!(data[root.get_child(Side::Left).unwrap()].value, 1);
251    }
252
253    #[test]
254    fn get_sibling_test() {
255        let rc = make_fake_tree_node();
256        let data = rc.borrow_mut();
257        assert_eq!(data[2].get_sibling(), Some(4));
258        assert_eq!(data[4].get_sibling(), Some(2));
259        assert_eq!(data[0].get_sibling(), None);
260        assert_eq!(data[3].get_sibling(), Some(1));
261        assert_eq!(data[1].get_sibling(), Some(3));
262        assert_eq!(data[5].get_sibling(), None);
263    }
264
265    #[test]
266    fn get_uncle_test() {
267        let rc = make_fake_tree_node();
268        let data = rc.borrow_mut();
269        assert_eq!(data[0].get_uncle(), None);
270        assert_eq!(data[1].get_uncle(), Some(2));
271        assert_eq!(data[2].get_uncle(), None);
272        assert_eq!(data[3].get_uncle(), Some(2));
273        assert_eq!(data[4].get_uncle(), None);
274        assert_eq!(data[5].get_uncle(), Some(4));
275    }
276
277    #[test]
278    fn get_size() {
279        let rc = make_fake_tree_node();
280        let data = rc.borrow_mut();
281        assert_eq!(data[0].get_size(), 6);
282        assert_eq!(data[4].get_size(), 3);
283        assert_eq!(data[5].get_size(), 1);
284    }
285
286    #[test]
287    fn get_height() {
288        let rc = make_fake_tree_node();
289        let data = rc.borrow_mut();
290        assert_eq!(data[0].get_height(), 3);
291        assert_eq!(data[4].get_height(), 2);
292        assert_eq!(data[5].get_height(), 1);
293    }
294
295    #[test]
296    fn find_min() {
297        let rc = make_fake_tree_node();
298        let data = rc.borrow_mut();
299        assert_eq!(data[(data[0].find_min())].value, 0);
300        assert_eq!(data[(data[2].find_min())].value, 6);
301    }
302
303    #[test]
304    fn find_leaf_count() {
305        let rc = make_fake_tree_node();
306        let data = rc.borrow_mut();
307        println!("{}", data[0].to_pretty_string(1));
308        assert_eq!(data[0].get_leaf_count(), 3);
309    }
310}
311