leetcode 863. 二叉树中所有距离为 K 的结点 -结点向上调整

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

来源：力扣（LeetCode）
链接：https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree

思路

• DFS 先找到target的子树中，相距target为k的节点
• 递归回溯的时候调整target节点为根节点（如果target已经是根节点，就不需要再调整了）
  • 删除target的子树
  • target的一个孩子置空，一个孩子指向父结点
  • 父节点指向target的指针改为指向父节点的父节点
    具体操作如图
• 调整后，target变成了root节点，直接递归找相距target为k的节点

如图

题解- 两次dfs- 通过 O(n)

# 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: int):
        if k==0:return [target.val]
        ans = []
        # 找到每个节点距离target的距离，如果等于K 则返回

        # 找到目标节点target
        # 找到目标节点相距K的节点
        def dfs(root,l,r,papa,ad = True):
            if not root: return False
            if l+r == k:
                ans.append(root.val)
                return False
            if root == target:l = r = k
            f1 = dfs(root.left,l-1,r,root)
            f2 = dfs(root.right,l,r-1,root)
            if ad:
                adjust(root,papa)
                if f1:root.left = papa
                if f2:root.right = papa
                if root==target:
                    root.left = None
                    root.right = papa
            if root == target or f1 or f2:return True
            return False

        def adjust(root,papa):
            if root == target:
                target.left == None
                target.right = papa
        if k == 0:return [target]
        if root!=target:dfs(root,-1,-1,None)
        dfs(target,-1,-1,None,ad=False)
        return ans