二叉树中所有距离为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