题目
寻找距离与目标的差距为K的各个值。
代码
# 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, target, K):
conn = collections.defaultdict(list)
def connect(parent, child):
if parent and child:
conn[parent.val].append(child.val)
conn[child.val].append(parent.val)
if child.left: connect(child, child.left)
if child.right: connect(child, child.right)
connect(None, root)
bfs = [target.val]
seen = set(bfs)
# 从5开始扩展2次。
for i in range(K):
bfs = [y for x in bfs for y in conn[x] if y not in seen]
seen |= set(bfs)
return bfs
高效代码
# 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, target, K):
"""
:type root: TreeNode
:type target: TreeNode
:type K: int
:rtype: List[int]
"""
def dfs(node,par=None):
if node:
node.par=par
dfs(node.left,node)
dfs(node.right,node)
dfs(root)
# 给树添加上父结点。
# 目标
seen={target}
queue=collections.deque([(target,0)])
# 往队列+ 上一层的从前面扔出去,然后把新的添加到后面。除去掉seen的。
while queue:
if queue[0][1]==K:
return [node.val for node,d in queue]
node,d=queue.popleft()
for nei in (node.left,node.right,node.par):
if nei and nei not in seen:
seen.add(nei)
queue.append((nei, d+1))
return []