[LeetCode/LintCode] Binary Tree Vertical Order Traversal

Problem

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.

Example

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]
]

Solution

class Solution {   
    public List<List<Integer>> verticalOrder(TreeNode root) {
        //result array and edge case
        List<List<Integer>> res = new ArrayList<>();
        if (root == null) return res;
        //create map and queues
        Map<Integer, List<Integer>> map = new HashMap<>();
        Queue<TreeNode> nodes = new LinkedList<>();
        Queue<Integer> indices = new LinkedList<>();
        //initialize the queues
        nodes.offer(root);
        indices.offer(0);
        //bfs to traverse all nodes
        while (!nodes.isEmpty()) {
            TreeNode node = nodes.poll();
            Integer index = indices.poll();
            map.putIfAbsent(index, new ArrayList<Integer>());
            map.get(index).add(node.val);
            
            if (node.left != null) {
                nodes.offer(node.left);
                indices.offer(index-1);
            }
            
            if (node.right != null) {
                nodes.offer(node.right);
                indices.offer(index+1);
            }
        }
        //well, the list needs to be ordered
        for (int i = Collections.min(map.keySet()); i <= Collections.max(map.keySet()); i++) {
            res.add(map.get(i));
        }
        return res;
    }
}