leetcode 987. Vertical Order Traversal of a Binary Tree

Question

anyone can provide some suggestion about my answer for the questions 987. Vertical Order Traversal of a Binary Tree

Given a binary tree, return the vertical order traversal of its nodes values. For each node at position (X, Y), its left and right children respectively will be at positions (X-1, Y-1) and (X+1, Y-1). Running a vertical line from X = -infinity to X = +infinity, whenever the vertical line touches some nodes, we report the values of the nodes in order from top to bottom (decreasing Y coordinates). If two nodes have the same position, then the value of the node that is reported first is the value that is smaller.

Return an list of non-empty reports in order of X coordinate. Every report will have a list of values of nodes.

Input: [3,9,20,null,null,15,7]

Given binary tree [1, 2, 3, 4, 6, 5, 7]

      1
    /   \
   2     3
  / \   / \
 4   6 5   7

return its vertical order traversal as

Output: [[4], [2], [1, 5, 6], [3], [7]]

class Solution:
    def verticalOrder(self, root: TreeNode) -> List[List]:
        stack = [(root, 0)]
        node_map = dict()
        while stack:
            tmp = dict()
            for _ in range(len(stack)):
                node, index = stack.pop(0)
                if node:
                    tmp[index] = tmp.get(index, []) + [node.val]
                if node.left: stack.append((node.left, index - 1))
                if node.right: stack.append((node.right, index + 1))

            tmp = {k: sorted(i) for k, i in tmp.items()}
            node_map = {k: node_map.get(k, []) + tmp.get(k, []) for k in list(tmp.keys()) + list(node_map.keys())}
        res = sorted([(index, val) for index, val in node_map.items()], key=lambda x: x[0])
        return [i for _, i in res]

Answer 1

Your code seems to have the right general structure, but is doing some extra work.

Comments would help people (including you in the future) understand what your code does and why (e.g the comment below explaining why -y is used).

Also, the problem says that nodes at the same x-coord are listed in order of their y-coord (top-to-bottom) and then by value, but you don't seem to keep track of the y-coord.

from collections import defaultdict

class Solution:
    def verticalOrder(self, root: TreeNode) -> List[List]:
        # the stack keeps track of nodes that are waiting to be
        # processed, along with their x and y coordinates
        stack = [(root, 0, 0)]

        # the key is x-coordinate of a tree node, the value is a list of (-y, node_value) tuples.
        # We use -y so so that when the list is sorted later, the tuples are sorted top-to-bottom
        # and then by node_value
        node_map = defaultdict(list)

        # while there are nodes left to be processed, get the next one and add it to the node_map
        while stack:
            node, x, y = stack.pop()
            node_map[x].append((-y, node.val))

            # add the node's children to the stack for processing
            if node.left: 
                stack.append((node.left, x - 1, y - 1))

            if node.right:
                stack.append((node.right, x + 1, y - 1))

        result = []

        # process groups of nodes in order from smallest to largest x-coordinate
        for x_coord in sorted(node_map.keys()):

            # sort the tuples by y-coordinate and then value
            values_at_this_x = [node_value for _,node_value in sorted(node_map[x_coord])]

            result.append(values_at_this_x)

        return result