二叉树中所有距离为 K 的结点

标签: 深度优先搜索 广度优先搜索 二叉树

难度: Medium

给定一个二叉树(具有根结点 root), 一个目标结点 target ,和一个整数值 k

返回到目标结点 target 距离为 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

示例 2:

输入: root = [1], target = 1, k = 3
输出: []

提示:

  • 节点数在 [1, 500] 范围内
  • 0 <= Node.val <= 500
  • Node.val 中所有值 不同
  • 目标结点 target 是树上的结点。
  • 0 <= k <= 1000

Submission

运行时间: 28 ms

内存: 0.0 MB

# 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]:
#         def dfs(node, par = None):
#             if node:
#                 node.par = par
#                 dfs(node.left, node)
#                 dfs(node.right, node)

#         dfs(root)
        
#         queue = collections.deque([(target,0)])
#         seen = {target}
        
#         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 []
        
        ans = []

        # Return distance from node to target if exists, else -1
        # Vertex distance: the # of vertices on the path from node to target
        def dfs(node):
            if not node:
                return -1
            elif node is target:
                subtree_add(node, 0)
                return 1
            else:
                L, R = dfs(node.left), dfs(node.right)
                if L != -1:
                    if L == k: ans.append(node.val)
                    subtree_add(node.right, L + 1)
                    return L + 1
                elif R != -1:
                    if R == k: ans.append(node.val)
                    subtree_add(node.left, R + 1)
                    return R + 1
                else:
                    return -1

        # Add all nodes 'K - dist' from the node to answer.
        def subtree_add(node, dist):
            if not node:
                return
            elif dist == k:
                ans.append(node.val)
            else:
                subtree_add(node.left, dist + 1)
                subtree_add(node.right, dist + 1)

        dfs(root)
        return ans
            

Explain

这个题解使用了两个DFS遍历来解决问题。第一个DFS遍历从根节点开始,用于查找目标节点target并记录从根节点到目标节点的路径长度。在找到目标节点时,调用第二个DFS遍历subtree_add,将距离目标节点距离为k的所有节点值加入结果列表。在第一个DFS遍历中,如果当前节点到目标节点的距离为k,则将当前节点值加入结果列表;然后继续递归遍历另一个子树,并将距离加1。最后返回结果列表。

时间复杂度: O(N)

空间复杂度: O(N)

class Solution:
    def distanceK(self, root: TreeNode, target: TreeNode, k: int) -> List[int]:
        ans = []

        # Return distance from node to target if exists, else -1
        # Vertex distance: the # of vertices on the path from node to target
        def dfs(node):
            if not node:
                return -1
            elif node is target:
                subtree_add(node, 0)
                return 1
            else:
                L, R = dfs(node.left), dfs(node.right)
                if L != -1:
                    if L == k: ans.append(node.val)
                    subtree_add(node.right, L + 1)
                    return L + 1
                elif R != -1:
                    if R == k: ans.append(node.val)
                    subtree_add(node.left, R + 1)
                    return R + 1
                else:
                    return -1

        # Add all nodes 'K - dist' from the node to answer.
        def subtree_add(node, dist):
            if not node:
                return
            elif dist == k:
                ans.append(node.val)
            else:
                subtree_add(node.left, dist + 1)
                subtree_add(node.right, dist + 1)

        dfs(root)
        return ans

Explore

在此题解中,确保不重复添加节点主要依靠的是控制递归过程的路径。当通过`dfs`函数找到目标节点`target`后,对其子树调用`subtree_add`函数是从目标节点开始,沿着不同的分支向下递归。此时,其他分支的节点不会被再次访问,因为`dfs`函数在找到目标节点并进行处理后,就不会再继续对已访问的分支进行递归。这样,每个节点最多被访问一次,不会被重复添加到结果列表中。

在二叉树中,一个节点在其左右子树中只能找到一次。如果左子树中找到了目标节点,则右子树中不可能再次找到同一个目标节点,反之亦然。因此,题解中的逻辑是符合二叉树的特性的。当`dfs`函数在左或右子树中找到目标节点后,它会停止进一步搜索其他路径并返回一个非-1的值,表明已经找到了目标节点。这保证了一旦目标节点被找到,不会再对其他分支进行无谓的搜索。

是的,`Vertex distance`指的是从起点到终点的路径中所包含的顶点数。这通常包括路径的起点和终点。例如,从节点A直接到节点B的顶点距离是2。这种定义在二叉树问题中常用来精确描述从一个节点到另一个节点所需要经过的节点总数。

题解中没有显式处理`k`大于树的最大高度的情况。在这种情况下,`subtree_add`函数将会尝试递归到比实际树的高度还要深的层级,但由于没有更多的节点,递归会在检查`node`是否为`null`时停止。这确保了即使`k`的值过大,函数也不会出错,只是会进行一些不必要的递归调用。在实际应用中,可以通过先计算树的最大高度来优化这个过程,避免不必要的递归。