LeetCode 987. Vertical Order Traversal of a Binary Tree

Problem Statement

use std::rc::Rc;
use std::cell::RefCell;
use std::collections::{BTreeMap, VecDeque};
use std::borrow::Borrow;

impl Solution {
    pub fn vertical_traversal(root: Option<Rc<RefCell<TreeNode>>>) -> Vec<Vec<i32>> {
        let mut map: BTreeMap<i32, Vec<(i32, i32)>> = BTreeMap::new();
        if root.is_none() {
            return vec![];
        }
        let mut q = VecDeque::new();
        q.push_back((root.unwrap(), (0, 0)));

        while !q.is_empty() {
            let cur = q.front().unwrap().clone();
            q.pop_front();
            let node = cur.1;
            let (h, v) = (node.0, node.1);
            map.entry(h).or_insert(vec![]);
            map.get_mut(&h).unwrap().push((v, RefCell::borrow(cur.0.borrow()).val));

            let (left, right) = (
                RefCell::borrow(&cur.0).left.clone(),
                RefCell::borrow(&cur.0).right.clone());

            if !left.is_none() {
                q.push_back((left.unwrap().clone(), (h - 1, v + 1)));
            }
            if !right.is_none() {
                q.push_back((right.unwrap().clone(), (h + 1, v + 1)));
            }
        }

        let mut ans: Vec<Vec<i32>> = vec![vec![]; map.len()];
        let mut index = 0;
        for (_k, mut v) in map {
            v.sort();
            for val in v {
                ans[index].push(val.1);
            }
            index += 1;
        }
        return ans;
    }
}