863. All Nodes Distance K in Binary Tree

tree bfs dfs

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer 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.

Example 2:

Input: root = [1], target = 1, k = 3
Output: []

Constraints:

• The number of nodes in the tree is in the range [1, 500].

• 0 <= Node.val <= 500

• All the values Node.val are unique.

• target is the value of one of the nodes in the tree.

• 0 <= k <= 1000

分析

1. 建立父节点映射：由于二叉树节点没有指向父节点的指针，我们首先需要遍历整棵树，记录每个节点的父节点，以便后续能够向上遍历。

2. 广度优先搜索（BFS）：从目标节点target开始进行BFS，逐层向外扩展。每次处理当前层的节点时，检查其左子节点、右子节点和父节点，确保每个节点只访问一次。当距离达到k时，收集当前层的所有节点。

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


class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        parents = {}
        def dfs(node, parent):
            if node:
                parents[node] = parent
                dfs(node.left, node)
                dfs(node.right, node)
        dfs(root, None)
        q = deque([(target, 0)])
        visited = set([target])
        res = []
        while q:
            node, distance = q.popleft()
            if distance == k:
                res.append(node.val)
                continue #此路可停
            for neighbor in [parents[node], node.left, node.right]:
                if neighbor and neighbor not in visited:
                    q.append((neighbor, distance + 1))
                    visited.add(neighbor)
        return res