leetcode_rust/
vertical_order_traversal_of_a_binary_tree.rs

1#![allow(dead_code)]
2
3use std::cell::RefCell;
4use std::rc::Rc;
5
6use crate::binary_tree::TreeNode;
7
8struct Solution;
9
10// local success, remote failure
11impl Solution {
12    pub fn vertical_traversal(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
13        use std::iter;
14
15        let mut left: i32 = 0;
16        let mut right: i32 = 0;
17        preorder(root.clone(), 0, 0, &mut left, &mut right);
18        let mut res = iter::repeat(vec![])
19            .take((right + 1 - left) as usize)
20            .collect::<Vec<Vec<i32>>>();
21
22        preorder2(root, &mut res, 0, 0, left.clone());
23
24        res
25    }
26}
27
28fn preorder(
29    root: Option<Rc<RefCell<TreeNode>>>,
30    mut x: i32,
31    y: i32,
32    left: &mut i32,
33    right: &mut i32,
34) {
35    match root {
36        None => {}
37        Some(node) => {
38            *left = *left.min(&mut x);
39            *right = *right.max(&mut x);
40            preorder(node.borrow().left.clone(), x - 1, y - 1, left, right);
41            preorder(node.borrow().right.clone(), x + 1, y - 1, left, right);
42        }
43    }
44}
45
46fn preorder2(
47    root: Option<Rc<RefCell<TreeNode>>>,
48    mut res: &mut Vec<Vec<i32>>,
49    x: i32,
50    y: i32,
51    left: i32,
52) {
53    match root {
54        None => {}
55        Some(node) => {
56            res[(x - left) as usize].push(node.borrow().val);
57            preorder2(node.borrow().left.clone(), &mut res, x - 1, y - 1, left);
58            preorder2(node.borrow().right.clone(), &mut res, x + 1, y - 1, left);
59        }
60    }
61}
62
63#[cfg(test)]
64mod tests {
65    use super::*;
66
67    #[test]
68    fn test1() {
69        let tree = Some(Rc::new(RefCell::new(TreeNode {
70            val: 0,
71            left: None,
72            right: Some(Rc::new(RefCell::new(TreeNode {
73                val: 1,
74                left: Some(Rc::new(RefCell::new(TreeNode {
75                    val: 2,
76                    left: None,
77                    right: None,
78                }))),
79                right: None,
80            }))),
81        })));
82
83        assert_eq!(
84            Solution::vertical_traversal(tree),
85            vec![vec![0, 2], vec![1]]
86        );
87    }
88}
leetcode_rust/vertical_order_traversal_of_a_binary_tree.rs

leetcode_rust/
vertical_order_traversal_of_a_binary_tree.rs
