leetcode 863. 二叉树中所有距离为 K 的结点(dfs+递归)

• 题目描述

给定一个二叉树（具有根结点 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



注意，输入的 "root" 和 "target" 实际上是树上的结点。
上面的输入仅仅是对这些对象进行了序列化描述。

• 解法：dfs+递归

一开始自己的想法是用递归计算两个节点的距离，然后找到target节点，然后再找到与target节点距离为K的节点，但实现起来在如何在递归中找target节点卡住了。

看了官方题解，这道题需要用两次递归来做。同样通过计算节点的距离来找到目标的节点。

如果target节点在root节点的左子树中，且target节点深度为n，那么所有root节点右子树深度为K-n的节点到target的距离则为K。

算法思路：

用dfs遍历所有节点，返回值是node到target的距离。遍历过程距离的计算会出现以下4种情况：

1. node == target，把左子树距离target节点的距离为K的所有节点放入res数组（用于存结果）
2. target在node左子树，那么target距离node的距离为left + 1，然后找出target到右子树距离为K-(left + 1)的节点假如res中。（这里需要注意直接递归subtree(node.right, left+1)是在距离为left+1的基础上继续寻找left+ 1+右子树距离和为K的右子树中某个节点）。
3. target在node右子树，与第2种情况类似
4. 如果target不在node的左右子树中，那么直接返回-1，不做任何处理。

其中，subtree()函数是用来查找子树距离节点node为K的节点，并且存入res中。

# 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]:
        res = []
        def dfs(node):
            if not node:
                return -1
            elif node is target:
                subtree(node, 0)
                return 1
            else:
                left, right = dfs(node.left), dfs(node.right)
                if left != -1:
                    if left == K:
                        res.append(node.val)
                    subtree(node.right, left + 1)
                    return left+1
                elif right != -1:
                    if right == K:
                        res.append(node.val)
                    subtree(node.left, right + 1)
                    return right + 1
                else:
                    return -1
        def subtree(node, distance):
            if not node:
                return
            elif distance == K:
                res.append(node.val)
            else:
                subtree(node.left, distance + 1)
                subtree(node.right, distance + 1)
        dfs(root)
        return res

参考链接：https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree/solution/er-cha-shu-zhong-suo-you-ju-chi-wei-k-de-jie-dian-/