863. All Nodes Distance K in Binary Tree.md

863. All Nodes Distance K in Binary Tree

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node. The answer can be returned in any order.

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

Output: [7,4,1]

Explanation: 
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.



Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

思路

• 把二叉树当成 Graph，然后用 BFS 顺着 Target 查找第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: TreeNode, target: TreeNode, K: int) -> List[int]:
        # first, we create a graph to represent binary tree
        g = defaultdict(list)
        
        def buildGraph(parent, child):
            if parent:
                g[parent].append(child)
                g[child].append(parent)
            if child.left:
                buildGraph(child, child.left)
            if child.right:
                buildGraph(child, child.right)
        
        buildGraph(None, root)
        queue = deque()
        res = []
        seen = set()
        seen.add(target)
        queue.append(target)
        k = 0
        
        while queue and k <= K:
            size = len(queue)
            while size > 0:
                node = queue.popleft()
                if k == K:
                    res.append(node.val)
                for child in g[node]:
                    if child in seen:
                        continue
                    queue.append(child)
                    seen.add(child)
                size -= 1
            k += 1
        
        return res