cargo_leet/leetcode_env/
tree.rs

1//! Leetcode Tree Support
2
3use std::{
4    cell::RefCell,
5    collections::VecDeque,
6    fmt::{Debug, Formatter},
7    rc::Rc,
8};
9
10///Definition for a binary tree node.
11#[derive(PartialEq, Eq)]
12pub struct TreeNode {
13    /// The value stored at this node
14    pub val: i32,
15    /// Link to the left child if one exists
16    pub left: Option<Rc<RefCell<TreeNode>>>,
17    /// Link to the right child if one exists
18    pub right: Option<Rc<RefCell<TreeNode>>>,
19}
20
21impl TreeNode {
22    #[inline]
23    /// Creates a new [TreeNode] with no children and the value passed
24    pub fn new(val: i32) -> Self {
25        TreeNode {
26            val,
27            left: None,
28            right: None,
29        }
30    }
31
32    fn wrapped_node_maybe(val: Option<i32>) -> Option<Rc<RefCell<Self>>> {
33        val.map(|x| Rc::new(RefCell::new(Self::new(x))))
34    }
35}
36
37/// Wrapper class to make handling empty trees easier and building of trees
38/// easier via [From] impls
39#[derive(PartialEq, Eq)]
40pub struct TreeRoot {
41    /// The root of the tree held
42    pub root: Option<Rc<RefCell<TreeNode>>>,
43}
44
45impl Debug for TreeRoot {
46    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
47        let mut vec: Vec<Option<i32>> = self.into();
48
49        // Remove trailing None's
50        while !vec.is_empty() && vec[vec.len() - 1].is_none() {
51            let _ = vec.remove(vec.len() - 1);
52        }
53
54        let vec: Vec<String> = vec
55            .iter()
56            .map(|x| {
57                if let Some(x) = x {
58                    format!("{x}")
59                } else {
60                    "None".to_string()
61                }
62            })
63            .collect();
64        write!(f, "{vec:?}")
65    }
66}
67
68impl From<&TreeRoot> for Vec<Option<i32>> {
69    fn from(value: &TreeRoot) -> Self {
70        let mut result = vec![];
71        let mut deque = VecDeque::new();
72        if value.root.is_some() {
73            deque.push_back(value.root.clone());
74        }
75        while !deque.is_empty() {
76            let node = deque.pop_front().unwrap();
77            match &node {
78                Some(node) => {
79                    let node = node.as_ref().borrow();
80                    result.push(Some(node.val));
81                    deque.push_back(node.left.clone());
82                    deque.push_back(node.right.clone());
83                }
84                None => {
85                    result.push(None);
86                }
87            }
88        }
89
90        // Trim trailing None
91        while let Some(_last) = result.last() {
92            if _last.is_none() {
93                result.pop();
94            } else {
95                break;
96            }
97        }
98        result
99    }
100}
101
102impl From<Option<Rc<RefCell<TreeNode>>>> for TreeRoot {
103    fn from(root: Option<Rc<RefCell<TreeNode>>>) -> Self {
104        Self { root }
105    }
106}
107
108impl From<&str> for TreeRoot {
109    /// Expects the "[]" around the values, separated by comma "," and only
110    /// integers and "null" (which is the format you'll get on LeetCode)
111    ///
112    /// # Panics
113    ///
114    /// This function panics if it doesn't match the expected format
115    fn from(value: &str) -> Self {
116        let mut result: Vec<Option<i32>>;
117        // Check and remove "[]"
118        assert!(value.len() >= 2);
119        assert_eq!('[', value.chars().next().unwrap());
120        assert_eq!(']', value.chars().nth(value.len() - 1).unwrap());
121        if value.len() == 2 {
122            // Empty array return empty tree
123            return Self { root: None };
124        }
125        let value = &value[1..value.len() - 1];
126
127        // Separate by comma
128        let values: Vec<&str> = value.split(',').map(|v| v.trim()).collect();
129
130        // Convert into values
131        result = vec![];
132        for value in values {
133            result.push(if value == "null" {
134                None
135            } else {
136                Some(value.parse().unwrap())
137            })
138        }
139
140        result.into()
141    }
142}
143
144impl Debug for TreeNode {
145    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
146        let left = if let Some(left) = &self.left {
147            format!("{:?}", left.as_ref().borrow())
148        } else {
149            "None".to_string()
150        };
151        let right = if let Some(right) = &self.right {
152            format!("{:?}", right.as_ref().borrow())
153        } else {
154            "None".to_string()
155        };
156        write!(f, "{{val:{} left:{} right:{}}}", self.val, left, right)
157    }
158}
159
160impl From<Vec<i32>> for TreeRoot {
161    fn from(value: Vec<i32>) -> Self {
162        value
163            .iter()
164            .map(|x| Some(*x))
165            .collect::<Vec<Option<i32>>>()
166            .into()
167    }
168}
169
170impl From<Vec<Option<i32>>> for TreeRoot {
171    /// Converts the incoming vec into a tree
172    fn from(list: Vec<Option<i32>>) -> Self {
173        // Based on https://leetcode.com/problems/recover-binary-search-tree/solutions/32539/Tree-Deserializer-and-Visualizer-for-Python/
174
175        if list.is_empty() {
176            return TreeRoot { root: None };
177        }
178
179        let nodes: Vec<Option<Rc<RefCell<TreeNode>>>> = list
180            .iter()
181            .map(|x| TreeNode::wrapped_node_maybe(*x))
182            .collect();
183
184        let mut kids: Vec<Option<Rc<RefCell<TreeNode>>>> = nodes
185            .iter()
186            .rev()
187            .map(|x| x.as_ref().map(Rc::clone))
188            .collect();
189
190        let root = kids.pop().expect("Check for empty done at top");
191
192        for node in nodes.into_iter().flatten() {
193            if let Some(left_child) = kids.pop() {
194                node.borrow_mut().left = left_child;
195            }
196            if let Some(right_child) = kids.pop() {
197                node.borrow_mut().right = right_child;
198            }
199        }
200
201        TreeRoot { root }
202    }
203}
204
205impl From<TreeRoot> for Option<Rc<RefCell<TreeNode>>> {
206    fn from(value: TreeRoot) -> Self {
207        value.root
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    /// Creates the test tree seen below
216    /// Leetcode rep: [1,2,5,3,null,6,7,null,4,null,null,8]
217    ///            1
218    ///         /     \
219    ///        /       \
220    ///       /         \
221    ///      2           5
222    ///    /   \       /   \
223    ///   3     -     6     7
224    ///  / \         / \   / \
225    /// - 4       -   - 8   -
226    #[allow(unused_mut)] // It's easier to read the code if they all line up but the leaves  don't need
227                         // to be mutable
228    fn test_tree() -> Option<Rc<RefCell<TreeNode>>> {
229        let mut node1 = Some(Rc::new(RefCell::new(TreeNode::new(1))));
230        let mut node2 = Some(Rc::new(RefCell::new(TreeNode::new(2))));
231        let mut node3 = Some(Rc::new(RefCell::new(TreeNode::new(3))));
232        let mut node4 = Some(Rc::new(RefCell::new(TreeNode::new(4))));
233        let mut node5 = Some(Rc::new(RefCell::new(TreeNode::new(5))));
234        let mut node6 = Some(Rc::new(RefCell::new(TreeNode::new(6))));
235        let mut node7 = Some(Rc::new(RefCell::new(TreeNode::new(7))));
236        let mut node8 = Some(Rc::new(RefCell::new(TreeNode::new(8))));
237        node3.as_mut().unwrap().borrow_mut().right = node4;
238        node7.as_mut().unwrap().borrow_mut().left = node8;
239        node2.as_mut().unwrap().borrow_mut().left = node3;
240        node5.as_mut().unwrap().borrow_mut().left = node6;
241        node5.as_mut().unwrap().borrow_mut().right = node7;
242        node1.as_mut().unwrap().borrow_mut().left = node2;
243        node1.as_mut().unwrap().borrow_mut().right = node5;
244        node1
245    }
246
247    #[test]
248    fn from_tree_to_vec() {
249        // Arrange
250        let start: TreeRoot = test_tree().into();
251        let expected = vec![
252            Some(1),
253            Some(2),
254            Some(5),
255            Some(3),
256            None,
257            Some(6),
258            Some(7),
259            None,
260            Some(4),
261            None,
262            None,
263            Some(8),
264        ];
265
266        // Act
267        let actual: Vec<Option<i32>> = (&start).into();
268
269        // Assert
270        assert_eq!(actual, expected);
271    }
272
273    #[test]
274    fn from_str_to_tree() {
275        // Arrange
276        let start = "[1,2,5,3,null,6,7,null,4,null,null,8]";
277        let expected = test_tree();
278
279        // Act
280        let root: TreeRoot = start.into();
281        let actual: Option<Rc<RefCell<TreeNode>>> = root.into();
282
283        // Assert
284        assert_eq!(actual, expected);
285    }
286}
cargo_leet/leetcode_env/tree.rs

cargo_leet/leetcode_env/
tree.rs
