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

cargo_leet/leetcode_env/
tree.rs
