[leetcode] 987. Vertical Order Traversal of a Binary Tree @ python

原题

https://leetcode.com/problems/vertical-order-traversal-of-a-binary-tree/

解法

BFS. 由于题意要求先按照x顺序排列, 再按照y倒序排列, 最后按照节点的值顺序排列, 我们构建字典d, 格式为d[x] = [(y, node.val)], 用q储存每层的(node, x, y), 将每层的信息放入字典d. 然后对d的键进行升序排列, 得到每层的列表, 然后对列表先按照y倒序排列, 再按照node.val顺序排列, 最后将每层的node.val放入ans即可.
Time: O(n)
Space: O(n)

代码

# Definition for a binary tree node.
# class TreeNode(object):
#     def __init__(self, x):
#         self.val = x
#         self.left = None
#         self.right = None

class Solution(object):
    def verticalTraversal(self, root):
        """
        :type root: TreeNode
        :rtype: List[List[int]]
        """
        # use a dict to store info: d[x] = [(y, node.val)]
        d = collections.defaultdict(list)
        q = [(root, 0, 0)]
        ans = []
        while q:
            for i in range(len(q)):
                node, x, y = q.pop(0)
                d[x].append((y, node.val))
                if node.left:
                    q.append((node.left, x-1, y-1))
                if node.right:
                    q.append((node.right, x+1, y-1))
                    
        # sort the d first by x, then by y
        for x in sorted(d.keys()):
            level = [tup[1] for tup in sorted(d[x], key = lambda t:(-t[0], t[1]))]
            ans.append(level)
        return ans