cargo_leet/leetcode_env/
tree.rs

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

cargo_leet/leetcode_env/
tree.rs
