863. All Nodes Distance K in Binary Tree - Medium

We are given a binary tree (with root node root), a target node, and an integer value K.

Return a list of the values of all nodes that have a distance K from the target node.  The answer can be returned in any order.

 

Example 1:

Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

Output: [7,4,1]

Explanation: 
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.



Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.

 

Note:

1. The given tree is non-empty.
2. Each node in the tree has unique values 0 <= node.val <= 500.
3. The target node is a node in the tree.
4. 0 <= K <= 1000.

 

dis函数：从root出发，计算root到target的距离，如果root为null说明该树里没有target，返回-1；如果root = target，调用collect函数收集节点并返回0；如果都不是，递归调用dis函数计算左右左右子节点到target的距离l, r，如果l非负说明target在左子树中，如果l正好等于k - 1说明该节点到target的距离为k，放入res中，然后对右子树调用collect函数收集节点，此时传入的距离参数为k-l-2，最后返回l+1。对r的处理类似。如果l r 都为负，说明树里没有target，返回-1

time: O(n), space: O(n)

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        List<Integer> res = new ArrayList<>();
        dis(root, target, K, res);
        return res;
    }
    
    public int dis(TreeNode root, TreeNode target, int K, List<Integer> res) {
        if(root == null) {
            return -1;
        }
        if(root == target) {
            collect(root, K, res);
            return 0;
        }
        int l = dis(root.left, target, K, res);
        int r = dis(root.right, target, K, res);
        if(l >= 0) {
            if(l == K - 1) {
                res.add(root.val);
            }
            collect(root.right, K - l - 2, res);
            return l + 1;
        }
        if(r >= 0) {
            if(r == K - 1) {
                res.add(root.val);
            }
            collect(root.left, K - r - 2, res);
            return r + 1;
        }
        return -1;
    }
    
    public void collect(TreeNode node, int d, List<Integer> res) {
        if(node == null || d < 0) {
            return;
        }
        if(d == 0) {
            res.add(node.val);
        }
        collect(node.left, d - 1, res);
        collect(node.right, d - 1, res);
    }
}