314. Binary Tree Vertical Order Traversal

Given a binary tree, return the vertical order traversal of its nodes' values. (ie, from top to bottom, column by column).

If two nodes are in the same row and column, the order should be from left to right.

Examples 1:

Input: [3,9,20,null,null,15,7]

   3
  /\
 /  \
 9  20
    /\
   /  \
  15   7 

Output:

[
  [9],
  [3,15],
  [20],
  [7]
]

Examples 2:

Input: [3,9,8,4,0,1,7]

     3
    /\
   /  \
   9   8
  /\  /\
 /  \/  \
 4  01   7 

Output:

[
  [4],
  [9],
  [3,0,1],
  [8],
  [7]
]

Examples 3:

Input: [3,9,8,4,0,1,7,null,null,null,2,5] (0's right child is 2 and 1's left child is 5)

     3
    /\
   /  \
   9   8
  /\  /\
 /  \/  \
 4  01   7
    /\
   /  \
   5   2

Output:

[
  [4],
  [9,5],
  [3,0,1],
  [8,2],
  [7]
]

Rust Solution

struct Solution;
use rustgym_util::*;
use std::collections::BTreeMap;

trait Preorder {
    fn preorder(&self, row: i32, col: i32, grid: &mut BTreeMap<i32, BTreeMap<i32, Vec<i32>>>);
}

impl Preorder for TreeLink {
    fn preorder(&self, row: i32, col: i32, grid: &mut BTreeMap<i32, BTreeMap<i32, Vec<i32>>>) {
        if let Some(node) = self {
            let node = node.borrow();
            let val = node.val;
            let left = &node.left;
            let right = &node.right;
            grid.entry(col)
                .or_default()
                .entry(row)
                .or_default()
                .push(val);

            left.preorder(row + 1, col - 1, grid);
            right.preorder(row + 1, col + 1, grid);
        }
    }
}

impl Solution {
    fn vertical_order(root: TreeLink) -> Vec<Vec<i32>> {
        let mut grid: BTreeMap<i32, BTreeMap<i32, Vec<i32>>> = BTreeMap::new();
        root.preorder(0, 0, &mut grid);
        grid.into_iter()
            .map(|(_, row)| row.into_iter().flat_map(|(_, v)| v.into_iter()).collect())
            .collect()
    }
}

#[test]
fn test() {
    let root = tree!(3, tree!(9), tree!(20, tree!(15), tree!(7)));
    let res = vec_vec_i32![[9], [3, 15], [20], [7]];
    assert_eq!(Solution::vertical_order(root), res);
    let root = tree!(
        3,
        tree!(9, tree!(4), tree!(0)),
        tree!(8, tree!(1), tree!(7))
    );
    let res = vec_vec_i32![[4], [9], [3, 0, 1], [8], [7]];
    assert_eq!(Solution::vertical_order(root), res);
}

Having problems with this solution? Click here to submit an issue on github.