解法
树型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