力扣每日一题21.07.28

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

题目描述：

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

返回到目标结点距离为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"实际上是树上的结点。上面的输入仅仅是对这些对象进行了序列化描述。

提示：

• 给定的树是非空的。
• 树上的每个结点都具有唯一的值0 <= node.val <= 500。
• 目标结点target是树上的结点。
• 0 <= k <= 1000。

思路：

二叉树转图后BFS即可。

Python代码：

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

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        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)
        cur = [target.val]
        visited = {target.val}
        while k:
            next_time = []
            while cur:
                tmp = cur.pop()
                for node in graph[tmp]:
                    if node not in visited:
                        visited.add(node)
                        next_time.append(node)
            k -= 1
            cur = next_time
        return cur