leetgo_rs/
tree.rs

1use std::cell::RefCell;
2use std::collections::VecDeque;
3use std::rc::Rc;
4
5use serde::{Deserialize, Serialize, Serializer};
6use serde::de::SeqAccess;
7use serde::ser::SerializeSeq;
8
9// LeetCode use `Option<Rc<RefCell<TreeNode>>>` for tree links, but `Option<Box<TreeNode>>` should be enough.
10// https://github.com/pretzelhammer/rust-blog/blob/master/posts/learning-rust-in-2020.md#leetcode
11type TreeLink = Option<Rc<RefCell<TreeNode>>>;
12
13#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
14pub struct TreeNode {
15    pub val: i32,
16    pub left: TreeLink,
17    pub right: TreeLink,
18}
19
20impl TreeNode {
21  #[inline]
22  pub fn new(val: i32) -> Self {
23    TreeNode {
24      val,
25      left: None,
26      right: None
27    }
28  }
29}
30
31#[macro_export]
32macro_rules! tree {
33    () => {
34        None
35    };
36    ($e:expr) => {
37        Some(Rc::new(RefCell::new(TreeNode::new($e))))
38    };
39}
40
41#[derive(Debug, PartialEq, Eq, PartialOrd, Ord)]
42pub struct BinaryTree(TreeLink);
43
44impl From<BinaryTree> for TreeLink {
45    fn from(tree: BinaryTree) -> Self {
46        tree.0
47    }
48}
49
50impl From<TreeLink> for BinaryTree {
51    fn from(link: TreeLink) -> Self {
52        BinaryTree(link)
53    }
54}
55
56impl Serialize for BinaryTree {
57    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
58        where
59            S: Serializer,
60    {
61        let mut queue = VecDeque::new();
62        let mut nodes = Vec::new();
63        queue.push_back(self.0.clone());
64        while let Some(node) = queue.pop_front() {
65            nodes.push(node.clone());
66            if let Some(ref node) = node {
67                queue.push_back(node.borrow().left.clone());
68                queue.push_back(node.borrow().right.clone());
69            }
70        }
71        while nodes.len() > 0 && nodes.last().unwrap().is_none() {
72            nodes.pop();
73        }
74
75        let mut seq = serializer.serialize_seq(None)?;
76        for node in nodes {
77            if let Some(ref node) = node {
78                seq.serialize_element(&node.borrow().val.clone())?;
79            } else {
80                seq.serialize_element(&None::<i32>)?;
81            }
82        }
83        seq.end()
84    }
85}
86
87struct BinaryTreeVisitor;
88
89impl<'de> serde::de::Visitor<'de> for BinaryTreeVisitor {
90    type Value = BinaryTree;
91
92    fn expecting(&self, formatter: &mut std::fmt::Formatter) -> std::fmt::Result {
93        formatter.write_str("a list of optional integers")
94    }
95
96    fn visit_seq<A>(self, mut seq: A) -> Result<Self::Value, A::Error>
97        where
98            A: SeqAccess<'de>,
99    {
100        let mut nodes: Vec<TreeLink> = Vec::new();
101
102        while let Some(val) = seq.next_element::<Option<i32>>()? {
103            nodes.push(val.map(|v: i32| Rc::new(RefCell::new(TreeNode {
104                val: v,
105                left: None,
106                right: None,
107            }))));
108        }
109
110        let root = nodes[0].clone();
111        let (mut i, mut j) = (0, 1);
112
113        while j < nodes.len() {
114            if let Some(ref current_node) = nodes[i] {
115                current_node.borrow_mut().left = nodes[j].clone();
116                j += 1;
117                if j < nodes.len() {
118                    current_node.borrow_mut().right = nodes[j].clone();
119                    j += 1;
120                }
121            }
122            i += 1;
123        }
124
125        Ok(BinaryTree(root))
126    }
127}
128
129impl<'de> Deserialize<'de> for BinaryTree {
130    fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
131        where
132            D: serde::Deserializer<'de>,
133    {
134        deserializer.deserialize_seq(BinaryTreeVisitor)
135    }
136}
137
138#[cfg(test)]
139mod tests {
140    use super::*;
141
142    #[test]
143    fn test_tree_serialize() {
144        let root = TreeNode {
145            val: 1,
146            left: Some(Rc::new(RefCell::new(TreeNode {
147                val: 2,
148                left: None,
149                right: None,
150            }))),
151            right: Some(Rc::new(RefCell::new(TreeNode {
152                val: 4,
153                left: Some(Rc::new(RefCell::new(TreeNode {
154                    val: 3,
155                    left: None,
156                    right: None,
157                }))),
158                right: None,
159            }))),
160        };
161        let tree = BinaryTree(Some(Rc::new(RefCell::new(root))));
162        let serialized = serde_json::to_string(&tree).unwrap();
163        assert_eq!(serialized, "[1,2,4,null,null,3]");
164    }
165}
leetgo_rs/tree.rs

leetgo_rs/
tree.rs
