LeetCode: 987. Vertical Order Traversal of a Binary Tree

Solution

• Traverse freely. There's no way you could know how much the tree will stretch to left or right in advance. While you traverse, gather data in heap.
• If you go left, do column_depth--. If right, do column_depth++.
• You should consider row (level) as well, because there's a condition that if there are some nodes in the same column but different row, the node on the less row should come first.

Python Code

import heapq
class Solution:
    def verticalTraversal(self, root: Optional[TreeNode]) -> List[List[int]]:

        # traverse
        heap = []
        stack = [(0, 0, root)]
        while stack:
            col, row, node = stack.pop(0)

            heapq.heappush(heap, (col, row, node.val))
            if node.left:
                stack.append((col-1, row+1, node.left))
            if node.right:
                stack.append((col+1, row+1, node.right))

        # fetch
        vo = []
        o = []
        col = heap[0][0]
        while heap:
            c, r, val = heapq.heappop(heap)
            if col != c:
                col = c
                vo.append(o)
                o = []
            o.append(val)
        vo.append(o)
        return vo