Update Vertical Order Traversal Of a Binary Tree.java

varunu28 · web-flow · commit 089078d45578 · 2022-09-03T21:53:50.000-07:00
diff --git a/Hard/Vertical Order Traversal Of a Binary Tree.java b/Hard/Vertical Order Traversal Of a Binary Tree.java
@@ -15,42 +15,46 @@
  */
 class Solution {
   public List<List<Integer>> verticalTraversal(TreeNode root) {
-    Map<Integer, List<TreeCoordinate>> map = new TreeMap<>();
-    Queue<TreeCoordinate> queue = new LinkedList<>();
-    queue.add(new TreeCoordinate(root, 0, 0));
+    Map<Integer, List<NodePair>> map = new HashMap<>();
+    Queue<NodePair> queue = new LinkedList<>();
+    queue.add(new NodePair(0, 0, root));
+    int minColumn = 0;
+    int maxColumn = 0;
     while (!queue.isEmpty()) {
       int size = queue.size();
       while (size-- > 0) {
-        TreeCoordinate removed = queue.remove();
+        NodePair removed = queue.remove();
+        minColumn = Math.min(minColumn, removed.y);
+        maxColumn = Math.max(maxColumn, removed.y);
         map.computeIfAbsent(removed.y, k -> new ArrayList<>()).add(removed);
         if (removed.node.left != null) {
-          queue.add(new TreeCoordinate(removed.node.left, removed.x + 1, removed.y - 1));
+          queue.add(new NodePair(removed.x + 1, removed.y - 1, removed.node.left));
         }
         if (removed.node.right != null) {
-          queue.add(new TreeCoordinate(removed.node.right, removed.x + 1, removed.y + 1));
+          queue.add(new NodePair(removed.x + 1, removed.y + 1, removed.node.right));
         }
       }
     }
     List<List<Integer>> result = new ArrayList<>();
-    for (Integer key : map.keySet()) {
-      List<TreeCoordinate> temp = map.get(key);
-      Collections.sort(temp,
-          Comparator.comparingInt((TreeCoordinate o) -> o.x).thenComparingInt(o -> o.node.val));
-      result.add(temp.stream().map(a -> a.node.val).collect(Collectors.toList()));
+    for (int key = minColumn; key <= maxColumn; key++) {
+      List<NodePair> list = map.get(key);
+      result.add(list.stream()
+                 .sorted(Comparator.comparingInt((NodePair o) -> o.x).thenComparingInt(o -> o.node.val))
+                 .map(np -> np.node.val)
+                 .collect(Collectors.toList()));
     }
     return result;
   }
   
-  private static class TreeCoordinate {
+  class NodePair {
+    int x;
+    int y;
+    TreeNode node;
     
-    public TreeNode node;
-    public int x;
-    public int y;
-    
-    public TreeCoordinate(TreeNode node, int x, int y) {
-      this.node = node;
+    public NodePair(int x, int y, TreeNode node) {
       this.x = x;
       this.y = y;
+      this.node = node;
     }
   }
 }
