二叉树所有距离为K的节点

# 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):
            def find(node,dist):
                if node:
                    if dist == K:
                        res.append(node.val)
                    elif dist<K:
                        find(node.left,dist+1)
                        find(node.right,dist+1)
            
            if node:
                if node == target:
                    if K == 0:
                        res.append(node.val)
                        return 1
                    find(node.left,1)
                    find(node.right,1)
                    return 1 
                else:
                    left = dfs(node.left)
                    right = dfs(node.right)
                    if left<=0 and right<=0:
                        return 0
                    elif left > 0:
                        if left == K:
                            res.append(node.val)
                        else:
                            find(node.right,left+1)
                        return left+1
                    elif right > 0:
                        if right == K:
                            res.append(node.val)
                        else:
                            find(node.left,right+1)
                        return right +1
            return -1
        dfs(root)
        return res