863. All Nodes Distance K in Binary Tree

863. All Nodes Distance K in Binary Tree

先构建图

# 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 distanceK(self, root, target, K):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type K: int
        :rtype: List[int]
        """
        import collections
        mydic = collections.defaultdict(list)
        qu = [root]
        while qu:
            node = qu.pop(0)
            if node.left:
                mydic[node.val].append(node.left.val)
                mydic[node.left.val].append(node.val)
                qu.append(node.left)
            if node.right:
                mydic[node.val].append(node.right.val)
                mydic[node.right.val].append(node.val)
                qu.append(node.right)
        qu = [target.val]
        seen = set()
        while K > 0:
            K -= 1
            for i in range(len(qu)):
                node = qu.pop(0)
                seen.add(node)
                for j in mydic[node]:
                    if j not in seen:
                        qu.append(j)
        return qu