987. Vertical Order Traversal of a Binary Tree - cocoder39/coco39_LC GitHub Wiki

987. Vertical Order Traversal of a Binary Tree

similar to 314

class Solution:
    def verticalTraversal(self, root: Optional[TreeNode]) -> List[List[int]]:
        queue = collections.deque()
        if root:
            queue.append((root, 0, 0))
        
        column_map = collections.defaultdict(list)
        min_col, max_col = 0, 0
        while queue:
            node, row, col = queue.popleft()
            min_col = min(min_col, col)
            max_col = max(max_col, col)
            column_map[col].append((row, node.val))
            if node.left:
                queue.append((node.left, row+1, col-1))
            if node.right:
                queue.append((node.right, row+1, col+1))

        res = []
        for i in range(min_col, max_col+1):
            tmp = [val for row, val in sorted(column_map[i])]
            res.append(tmp)
        return res