rustgym/leetcode/
_814_binary_tree_pruning.rs

1struct Solution;
2use rustgym_util::*;
3
4trait Postorder {
5    fn postorder(self) -> Self;
6}
7
8impl Postorder for TreeLink {
9    fn postorder(self) -> Self {
10        if let Some(node) = self {
11            let val = node.borrow().val;
12            let left = node.borrow_mut().left.take();
13            let right = node.borrow_mut().right.take();
14            let left = left.postorder();
15            let right = right.postorder();
16            if left.is_none() && right.is_none() && val == 0 {
17                None
18            } else {
19                tree!(val, left, right)
20            }
21        } else {
22            None
23        }
24    }
25}
26
27impl Solution {
28    fn prune_tree(root: TreeLink) -> TreeLink {
29        root.postorder()
30    }
31}
32
33#[test]
34fn test() {
35    let root = tree!(1, None, tree!(0, tree!(0), tree!(1)));
36    let res = tree!(1, None, tree!(0, None, tree!(1)));
37    assert_eq!(Solution::prune_tree(root), res);
38    let root = tree!(
39        1,
40        tree!(0, tree!(0), tree!(0)),
41        tree!(1, tree!(0), tree!(1))
42    );
43    let res = tree!(1, None, tree!(1, None, tree!(1)));
44    assert_eq!(Solution::prune_tree(root), res);
45    let root = tree!(
46        1,
47        tree!(1, tree!(1, tree!(0), None), tree!(1)),
48        tree!(0, tree!(0), tree!(1))
49    );
50    let res = tree!(1, tree!(1, tree!(1), tree!(1)), tree!(0, None, tree!(1)));
51    assert_eq!(Solution::prune_tree(root), res);
52}
rustgym/leetcode/_814_binary_tree_pruning.rs

rustgym/leetcode/
_814_binary_tree_pruning.rs
