LeetCode - Binary Tree Vertical Order Traversal

BinaryTreeVerticalOrderTraversal.java

// 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:

// Given binary tree [3,9,20,null,null,15,7],
//    3
//  /\
//  /  \
//  9  20
//     /\
//    /  \
//  15   7
// return its vertical order traversal as:
// [
//  [9],
//  [3,15],
//  [20],
//  [7]
// ]
// Given binary tree [3,9,8,4,0,1,7],
//     3
//     /\
//    /  \
//    9   8
//  /\  /\
//  /  \/  \
//  4  01   7
// return its vertical order traversal as:
// [
//  [4],
//  [9],
//  [3,0,1],
//  [8],
//  [7]
// ]
// Given binary tree [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
// return its vertical order traversal as:
// [
//  [4],
//  [9,5],
//  [3,0,1],
//  [8,2],
//  [7]
// ]

private int offset = 0;
public List<List<Integer>> verticalOrder(TreeNode root) {
    List<List<Integer>> result = new ArrayList<List<Integer>>();
    List<Map<Integer, List<Integer>>> res = new ArrayList<Map<Integer, List<Integer>>>();
    if(root == null) return result;
    helper(res, root, 0, 0);
    for(Map<Integer, List<Integer>> map : res){
        List<Integer> lst = new ArrayList<Integer>();
        for(List<Integer> levelList : map.values()){
            lst.addAll(levelList);
        }
        result.add(lst);
    }
    return result;
}

private void helper(List<Map<Integer, List<Integer>>> res, TreeNode root, int col, int row){
    if(root == null) return;
    
    if(col + offset < 0){
        res.add(0, new TreeMap<Integer, List<Integer>>());
        offset++;
    } else if( col + offset == res.size()){
        res.add(new TreeMap<Integer, List<Integer>>());
    }
    
    Map<Integer, List<Integer>> map = res.get(col + offset);
    if(map.containsKey(row)){
        map.get(row).add(root.val);
    } else{
        List<Integer> lst = new ArrayList<Integer>();
        lst.add(root.val);
        map.put(row, lst);
    }
    
    helper(res, root.left, col - 1, row + 1);
    helper(res, root.right, col + 1, row + 1);
}

// https://leetcode.com/discuss/75054/5ms-java-clean-solution