863. All Nodes Distance K in Binary Tree（BFS学习）

题目

寻找距离与目标的差距为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 []