BFS之二叉树中所有距离为 K 的结点

LeetCode 863. 二叉树中所有距离为 K 的结点

给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 K 。

返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。

示例 1：

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
输出：[7,4,1]
解释：
所求结点为与目标结点（值为 5）距离为 2 的结点，
值分别为 7，4，以及 1

注意，输入的 “root” 和 “target” 实际上是树上的结点。
上面的输入仅仅是对这些对象进行了序列化描述。
先构图，然后进行广度优先遍历

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, K: int) -> List[int]:
        from collections import defaultdict
        # make graph
        graph = defaultdict(set)
        def dfs(root):
            if root.left:
                graph[root.val].add(root.left.val)
                graph[root.left.val].add(root.val)
                dfs(root.left)
            if root.right:
                graph[root.val].add(root.right.val)
                graph[root.right.val].add(root.val)
                dfs(root.right)
        dfs(root)
        # bfs
        curr = [target.val]
        visited = {target.val}
        while K:
            nex = []
            while curr:
                temp = curr.pop()
                for node in graph[temp]:
                    if node not in visited:
                        visited.add(node)
                        nex.append(node)
            K -= 1
            curr = nex
        return curr