cauly_rust_leetcode_utils/
binary_tree.rs

1use std::cell::RefCell;
2use std::rc::Rc;
3
4// #region LeetCode TreeNode
5#[derive(Debug, PartialEq, Eq)]
6pub struct TreeNode {
7    pub val: i32,
8    pub left: Option<Rc<RefCell<TreeNode>>>,
9    pub right: Option<Rc<RefCell<TreeNode>>>,
10}
11
12impl TreeNode {
13    #[inline]
14    pub fn new(val: i32) -> Self {
15        TreeNode {
16            val,
17            left: None,
18            right: None,
19        }
20    }
21}
22// endregion
23
24// #region Tree ext
25
26#[derive(PartialEq, Copy, Clone)]
27pub enum TraversalType {
28    Inorder,
29    Preorder,
30    Postorder,
31}
32
33pub trait RefTreeNode {
34    fn get_val(&self) -> Option<i32>;
35    fn get_left(&self) -> Option<Rc<RefCell<TreeNode>>>;
36    fn get_right(&self) -> Option<Rc<RefCell<TreeNode>>>;
37    fn set_val(&mut self, v: i32);
38    fn set_sub_nodes(
39        &self,
40        left: &Option<Rc<RefCell<TreeNode>>>,
41        right: &Option<Rc<RefCell<TreeNode>>>,
42    );
43    fn walk(&self, t: TraversalType, func: fn(Option<Rc<RefCell<TreeNode>>>));
44    fn walk_t<T>(
45        &self,
46        t: TraversalType,
47        result: &mut T,
48        func: fn(&mut T, Option<Rc<RefCell<TreeNode>>>),
49    );
50    fn aggregate<T>(
51        &self,
52        val_func: fn(Option<Rc<RefCell<TreeNode>>>) -> T,
53        aggr_func: fn(T, Option<T>, Option<T>) -> Option<T>,
54    ) -> Option<T>;
55    fn aggregate_t<R, T>(
56        &self,
57        result: &mut R,
58        val_func: fn(&mut R, Option<Rc<RefCell<TreeNode>>>) -> T,
59        aggr_func: fn(&mut R, T, Option<T>, Option<T>) -> Option<T>,
60    ) -> Option<T>;
61}
62
63impl RefTreeNode for Option<Rc<RefCell<TreeNode>>> {
64    fn get_val(&self) -> Option<i32> {
65        if let Some(n) = self {
66            let n = n.borrow();
67            Some(n.val)
68        } else {
69            None
70        }
71    }
72    fn get_left(&self) -> Option<Rc<RefCell<TreeNode>>> {
73        if let Some(n) = self {
74            let n = n.borrow();
75            if let Some(l) = &n.left {
76                return Some(Rc::clone(&l));
77            }
78        }
79        None
80    }
81    fn get_right(&self) -> Option<Rc<RefCell<TreeNode>>> {
82        if let Some(n) = self {
83            let n = n.borrow();
84            if let Some(l) = &n.right {
85                return Some(Rc::clone(&l));
86            }
87        }
88        None
89    }
90    fn set_val(&mut self, v: i32) {
91        if let Some(n) = self {
92            let mut n = n.borrow_mut();
93            n.val = v;
94        }
95    }
96    fn set_sub_nodes(
97        &self,
98        left: &Option<Rc<RefCell<TreeNode>>>,
99        right: &Option<Rc<RefCell<TreeNode>>>,
100    ) {
101        if let Some(n) = self {
102            let mut n = n.borrow_mut();
103            n.left = match left {
104                None => None,
105                Some(n) => Some(Rc::clone(n)),
106            };
107            n.right = match right {
108                None => None,
109                Some(n) => Some(Rc::clone(n)),
110            };
111        }
112    }
113    fn walk(&self, t: TraversalType, func: fn(Option<Rc<RefCell<TreeNode>>>)) {
114        if let Some(node) = self {
115            if t == TraversalType::Preorder {
116                func(Some(Rc::clone(&node)));
117            }
118            {
119                let n = node.borrow();
120                if let Some(left) = &n.left {
121                    Some(Rc::clone(left)).walk(t, func);
122                }
123            }
124            if t == TraversalType::Inorder {
125                func(Some(Rc::clone(&node)));
126            }
127            {
128                let n = node.borrow();
129                if let Some(right) = &n.right {
130                    Some(Rc::clone(right)).walk(t, func);
131                }
132            }
133            if t == TraversalType::Postorder {
134                func(Some(Rc::clone(&node)));
135            }
136        }
137    }
138    fn walk_t<T>(
139        &self,
140        t: TraversalType,
141        result: &mut T,
142        func: fn(&mut T, Option<Rc<RefCell<TreeNode>>>),
143    ) {
144        if let Some(node) = self {
145            if t == TraversalType::Preorder {
146                func(result, Some(Rc::clone(&node)));
147            }
148            {
149                let n = node.borrow();
150                if let Some(left) = &n.left {
151                    Some(Rc::clone(left)).walk_t(t, result, func);
152                }
153            }
154            if t == TraversalType::Inorder {
155                func(result, Some(Rc::clone(&node)));
156            }
157            {
158                let n = node.borrow();
159                if let Some(right) = &n.right {
160                    Some(Rc::clone(right)).walk_t(t, result, func);
161                }
162            }
163            if t == TraversalType::Postorder {
164                func(result, Some(Rc::clone(&node)));
165            }
166        }
167    }
168    fn aggregate<T>(
169        &self,
170        val_func: fn(Option<Rc<RefCell<TreeNode>>>) -> T,
171        aggr_func: fn(T, Option<T>, Option<T>) -> Option<T>,
172    ) -> Option<T> {
173        if let Some(node) = self {
174            let n = node.borrow();
175            let mut lval = None;
176            let mut rval = None;
177            if let Some(left) = &n.left {
178                lval = Some(Rc::clone(left)).aggregate(val_func, aggr_func);
179            }
180            if let Some(right) = &n.right {
181                rval = Some(Rc::clone(right)).aggregate(val_func, aggr_func);
182            }
183            let val = val_func(Some(Rc::clone(&node)));
184            aggr_func(val, lval, rval)
185        } else {
186            None
187        }
188    }
189    fn aggregate_t<R, T>(
190        &self,
191        result: &mut R,
192        val_func: fn(&mut R, Option<Rc<RefCell<TreeNode>>>) -> T,
193        aggr_func: fn(&mut R, T, Option<T>, Option<T>) -> Option<T>,
194    ) -> Option<T> {
195        if let Some(node) = self {
196            let n = node.borrow();
197            let mut lval = None;
198            let mut rval = None;
199            if let Some(left) = &n.left {
200                lval = Some(Rc::clone(left)).aggregate_t(result, val_func, aggr_func);
201            }
202            if let Some(right) = &n.right {
203                rval = Some(Rc::clone(right)).aggregate_t(result, val_func, aggr_func);
204            }
205            let val = val_func(result, Some(Rc::clone(&node)));
206            aggr_func(result, val, lval, rval)
207        } else {
208            None
209        }
210    }
211}
212
213impl TreeNode {
214    pub fn from_string(s: &str) -> Option<Rc<RefCell<TreeNode>>> {
215        let v: Vec<&str> = s.split(',').collect();
216        if v.len() == 0 {
217            None
218        } else {
219            let root = Rc::new(RefCell::new(TreeNode::new(v[0].parse().unwrap())));
220            let mut last_row = vec![Rc::clone(&root)];
221            let mut iter = v.iter().skip(1);
222            'outer: loop {
223                let mut new_row = Vec::new();
224                for parent in last_row.iter_mut() {
225                    let left = match iter.next() {
226                        Some(&"null") => None,
227                        Some(num) => {
228                            Some(Rc::from(RefCell::from(TreeNode::new(num.parse().unwrap()))))
229                        }
230                        _ => break 'outer,
231                    };
232                    let right = match iter.next() {
233                        Some(&"null") => None,
234                        Some(num) => {
235                            Some(Rc::from(RefCell::from(TreeNode::new(num.parse().unwrap()))))
236                        }
237                        _ => None,
238                    };
239                    Some(Rc::clone(parent)).set_sub_nodes(&left, &right);
240                    if left != None {
241                        new_row.push(left.unwrap());
242                    }
243                    if right != None {
244                        new_row.push(right.unwrap());
245                    }
246                }
247                last_row = new_row;
248                if last_row.len() == 0 {
249                    break;
250                }
251            }
252
253            Some(root)
254        }
255    }
256}
257// #endregion