【Leetcode】863. All Nodes Distance K in Binary Tree 863. 二叉树中所有距离为 K 的结点

解法

树型DP
solve(root)返回target距离root的父亲 有多远，如果返回-1表示target不在root里

• 假如我们得到了L=solve(root.left)，假如L!=-1，说明target在左子树，且距离root那么只要把右子树里距离root的距离是K-L的结点加入答案就行了，右边不需要再遍历了
• 假如我们得到了R=solve(root.right)，操作也是类似的

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

class Solution(object):
    def distanceK(self, root, target, K):
        """
        :type root: TreeNode
        :type target: TreeNode
        :type K: int
        :rtype: List[int]
        """
        ans = []
        def solve(root):
            if root==None:
                return -1
            if root == target:
                add_node(root,K)
                return 1
            L,R = solve(root.left), solve(root.right)
            if L!=-1:
                if L==K:
                    ans.append(root.val)
                else:
                    add_node(root.right, K-L-1)
                return L+1
            if R!=-1:
                if R==K:
                    ans.append(root.val)
                else:
                    add_node(root.left,K-R-1)
                return R+1
            return -1
        def add_node(root,L):
            if root==None:
                return
            if L==0:
                ans.append(root.val)
            else:
                add_node(root.left,L-1)
                add_node(root.right,L-1)
        solve(root)
        return ans