lcrt/
tree.rs

1use std::{cell::RefCell, rc::Rc};
2
3/// Definition for a binary tree node. The original definition is shown as below:
4///
5/// ```
6/// #[derive(Debug, PartialEq, Eq)]
7/// pub struct TreeNode {
8///   pub val: i32,
9///   pub left: Option<Rc<RefCell<TreeNode>>>,
10///   pub right: Option<Rc<RefCell<TreeNode>>>,
11/// }
12///
13/// impl TreeNode {
14///   pub fn new(val: i32) -> Self { /* ... */ }
15/// }
16/// ```
17#[derive(Debug, PartialEq, Eq)]
18pub struct TreeNode {
19    pub val: i32,
20    pub left: Option<Rc<RefCell<TreeNode>>>,
21    pub right: Option<Rc<RefCell<TreeNode>>>,
22}
23
24impl TreeNode {
25    pub fn new(val: i32) -> Self {
26        TreeNode {
27            val,
28            left: None,
29            right: None,
30        }
31    }
32
33    #[cfg(feature = "testing")]
34    fn with_rc(val: i32) -> Rc<RefCell<TreeNode>> {
35        Self::with_children(val, None, None)
36    }
37
38    #[cfg(feature = "testing")]
39    fn with_children(
40        val: i32,
41        left: Option<Rc<RefCell<TreeNode>>>,
42        right: Option<Rc<RefCell<TreeNode>>>,
43    ) -> Rc<RefCell<TreeNode>> {
44        Rc::new(RefCell::new(TreeNode { val, left, right }))
45    }
46
47    /// Factory method to generate a tree from string. If string is empty, we return None.
48    /// We could not use Into or TryInto traits here, because it doesn't return Option as we wanted here.
49    #[cfg(feature = "testing")]
50    pub fn from_str(raw_value: &str) -> Option<Rc<RefCell<TreeNode>>> {
51        use std::collections::VecDeque;
52
53        let values: Vec<Option<i32>> = raw_value
54            .split(',')
55            .filter(|&x| !x.trim().is_empty())
56            .map(|v| match v.trim() {
57                "null" => None,
58                v => Some(i32::from_str_radix(v, 10).unwrap()),
59            })
60            .collect();
61
62        if values.is_empty() {
63            return None;
64        }
65
66        let head = Self::with_rc(values[0].unwrap());
67        let mut queue = VecDeque::new();
68        queue.push_back(head.clone());
69
70        let mut i = 1usize;
71        while !queue.is_empty() {
72            let parent_node = queue.pop_front().unwrap();
73            if i < values.len() {
74                if let Some(v) = values[i] {
75                    let child = Self::with_rc(v);
76                    queue.push_back(child.clone());
77                    parent_node.borrow_mut().left = Some(child);
78                }
79                i += 1;
80            } else {
81                break;
82            }
83
84            if i < values.len() {
85                if let Some(v) = values[i] {
86                    let child = Self::with_rc(v);
87                    queue.push_back(child.clone());
88                    parent_node.borrow_mut().right = Some(child);
89                }
90                i += 1;
91            } else {
92                break;
93            }
94        }
95
96        Some(head)
97    }
98
99    #[cfg(feature = "testing")]
100    pub fn to_string(root: Option<Rc<RefCell<TreeNode>>>) -> String {
101        use std::collections::VecDeque;
102
103        let mut result = String::new();
104
105        let mut queue = VecDeque::new();
106        let mut non_null_node_in_queue_count = if root.is_some() { 1 } else { 0 };
107        queue.push_back(root);
108
109        while !queue.is_empty() {
110            let node = queue.pop_front().unwrap();
111            match node {
112                Some(node) => {
113                    non_null_node_in_queue_count -= 1;
114                    result.push_str(&format!("{},", node.borrow().val));
115
116                    queue.push_back(node.borrow().left.clone());
117                    non_null_node_in_queue_count +=
118                        if node.borrow().left.is_some() { 1 } else { 0 };
119
120                    queue.push_back(node.borrow().right.clone());
121                    non_null_node_in_queue_count +=
122                        if node.borrow().right.is_some() { 1 } else { 0 };
123                }
124                None => {
125                    result.push_str("null,");
126                }
127            }
128
129            if non_null_node_in_queue_count == 0 {
130                break;
131            }
132        }
133
134        result = result.trim_end_matches(',').to_string();
135        result
136    }
137
138    /// Asserts 2 trees are equal.
139    #[cfg(feature = "testing")]
140    pub fn assert_eq(t1: Option<Rc<RefCell<TreeNode>>>, t2: Option<Rc<RefCell<TreeNode>>>) {
141        use pretty_assertions::assert_eq;
142
143        let t1_str = Self::to_string(t1);
144        let t2_str = Self::to_string(t2);
145        assert_eq!(t1_str, t2_str);
146    }
147
148    /// Asserts 2 trees are equal.
149    #[cfg(feature = "testing")]
150    pub fn assert_list_eq(
151        t1: Vec<Option<Rc<RefCell<TreeNode>>>>,
152        t2: Vec<Option<Rc<RefCell<TreeNode>>>>,
153    ) {
154        use pretty_assertions::assert_eq;
155        use std::collections::BTreeSet;
156
157        let t1_set = t1
158            .iter()
159            .map(|v| Self::to_string(v.clone()))
160            .collect::<BTreeSet<String>>();
161
162        let t2_set = t2
163            .iter()
164            .map(|v| Self::to_string(v.clone()))
165            .collect::<BTreeSet<String>>();
166
167        assert_eq!(t1_set, t2_set);
168    }
169}
170
171/// Macro for generate tree with vec-like syntax. If no parameter is provided, it returns None.
172/// More on how Leetcode serializes the tree can be found here: https://support.leetcode.com/hc/en-us/articles/360011883654-What-does-1-null-2-3-mean-in-binary-tree-representation-
173///
174/// Example:
175/// ```
176/// lc_tree!(""); // This returns None.
177/// lc_tree!("1,2,3"); // This returns a list: 1->(2, 3)
178/// lc_tree!("1,null,2,3"); // This returns a list: 1->(null, 2->(3, null))
179/// ```
180#[cfg(feature = "testing")]
181#[macro_export]
182macro_rules! lc_tree {
183    ($te:literal) => {
184        TreeNode::from_str($te)
185    };
186}
187
188/// Macro for testing 2 trees are equal. The second parameter is provided in a way of string. If empty string is provided, it check against None.
189///
190/// Example:
191/// ```
192/// lc_tree_assert_eq(t, ""); // Check against None.
193/// lc_tree_assert_eq(t, "1,2,3"); // Check against list: 1->(2, 3).
194/// lc_tree_assert_eq(t, "1,null,2,3"); // Check against list: 1->(null, 2->(3, null)).
195/// ```
196#[cfg(feature = "testing")]
197#[macro_export]
198macro_rules! lc_tree_assert_eq {
199    ($t:expr, $te:literal) => {
200        TreeNode::assert_eq($t, TreeNode::from_str($te))
201    };
202}
203
204/// Macro for testing 2 list of trees are equal. The second parameter is provided in a way of list. If () is provided, it check against None.
205///
206/// Example:
207/// ```
208/// lc_tree_assert_list_eq(t, ()); // Check against None.
209/// lc_tree_assert_list_eq(t, ("1,2,3")); // Check against list: 1->(2, 3).
210/// lc_tree_assert_list_eq(t, ("1,2", "1,null,2,3")); // Check against list: 1->(2), 1->(null, 2->(3, null)).
211/// ```
212#[cfg(feature = "testing")]
213#[macro_export]
214macro_rules! lc_tree_assert_list_eq {
215    ($t:expr, ($($te:literal),*)) => {
216        TreeNode::assert_list_eq($t, vec![$(TreeNode::from_str($te)),*])
217    };
218}
219
220#[cfg(test)]
221mod tests {
222    use super::*;
223    use cool_asserts::assert_panics;
224    use pretty_assertions::assert_eq;
225
226    #[test]
227    fn lc_tree_macro_can_generate_tree_from_empty_str() {
228        let t = lc_tree!("");
229        assert!(t.is_none());
230    }
231
232    #[test]
233    fn lc_tree_macro_can_generate_tree_from_non_empty_str() {
234        let t: Option<Rc<RefCell<TreeNode>>> = lc_tree!("1,null,2,3");
235        assert!(t.is_some());
236
237        let t = t.as_ref().unwrap().borrow();
238        assert_eq!(t.val, 1);
239        assert!(t.left.is_none());
240        assert!(t.right.is_some());
241        assert_eq!(t.right.as_ref().unwrap().borrow().val, 2);
242
243        let trc = t.right.as_ref().unwrap().borrow();
244        assert!(trc.left.is_some());
245        assert_eq!(trc.left.as_ref().unwrap().borrow().val, 3);
246        assert!(trc.right.is_none());
247
248        let trc2 = trc.left.as_ref().unwrap().borrow();
249        assert_eq!(trc2.val, 3);
250        assert!(trc2.left.is_none());
251        assert!(trc2.right.is_none());
252    }
253
254    #[test]
255    fn tree_node_assert_eq_can_pass_on_identical_trees() {
256        lc_tree_assert_eq!(lc_tree!(""), "");
257        lc_tree_assert_eq!(lc_tree!("1"), "1");
258        lc_tree_assert_eq!(lc_tree!("1,2,3"), "1,2,3");
259        lc_tree_assert_eq!(lc_tree!("1,null,2,3"), "1,null,2,3");
260        lc_tree_assert_eq!(lc_tree!("1,2,null,3"), "1,2,null,3");
261    }
262
263    #[test]
264    fn tree_node_assert_eq_can_fail_on_non_identical_trees() {
265        assert_panics!(lc_tree_assert_eq!(lc_tree!(""), "1"));
266        assert_panics!(lc_tree_assert_eq!(lc_tree!("1"), ""));
267        assert_panics!(lc_tree_assert_eq!(lc_tree!("1,null,2"), "1,2,null"));
268        assert_panics!(lc_tree_assert_eq!(lc_tree!("1,null,2,3"), "1,2,null,3"));
269    }
270
271    #[test]
272    fn tree_node_assert_list_eq_can_pass_on_identical_trees() {
273        lc_tree_assert_list_eq!(vec![lc_tree!("")], (""));
274        lc_tree_assert_list_eq!(vec![lc_tree!("1")], ("1"));
275        lc_tree_assert_list_eq!(vec![lc_tree!("1"), lc_tree!("1,null,2")], ("1", "1,null,2"));
276    }
277}