863. All Nodes Distance K in Binary Tree [Medium]

https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/

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.

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */

---

算法思路：

方法1：We first do a depth first search where we annotate every node with information about it's parent.

After, we do a breadth first search to find all nodes a distance K from the target.

class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        unordered_map<TreeNode*, TreeNode*> parent;
        dfsAndCaculateParent(root, nullptr, parent);
        
        queue<TreeNode*> q;
        unordered_set<TreeNode*> seen;
        q.push(target);
        seen.insert(target);
        TreeNode* node;
        int dist = 0;
        vector<int> res;
        while(!q.empty()) {
            for (int i = 0, n = q.size(); i < n; i++) {
                node = q.front(); q.pop();
                if (dist == K) {
                    res.push_back(node->val);
                } else {
                    if (node->left && seen.find(node->left) == seen.end()) {
                        q.push(node->left);
                        seen.insert(node->left);
                    }
                    if (node->right && seen.find(node->right) == seen.end()) {
                        q.push(node->right);
                        seen.insert(node->right);
                    }
                    if (parent[node] && seen.find(parent[node]) == seen.end()) {
                        q.push(parent[node]);
                        seen.insert(parent[node]);
                    }
                }
            }
            if (dist++ == K) break;
        }
        
        return res;
    } 
private:
    void dfsAndCaculateParent(TreeNode* node, TreeNode* parentNode, 
                              unordered_map<TreeNode*, TreeNode*>& parent) {
        if (node == nullptr) return;
        parent[node] = parentNode;
        dfsAndCaculateParent(node->left, node, parent);
        dfsAndCaculateParent(node->right, node, parent);
    }
};

 

方法2：比较麻烦，dfs，当找到target的时候，使用subtree_add到以target为根所在节点的子树中去寻找到案，并且返回1，这样节点target的父节点node中L = 1（假设是node的左子节点，不失一般性，后边调用情况对称变换即可），所以还需要到node的右节点所在子树中去找是否存在K - (L + 1)距离的节点，因为target是node的左子节点，所以到达node右子节点的距离是2，也就是(L + 1)，所以从node->right点找的时候距离已经是L + 1，也就是从node->right为根的子树中去找距离为K - (L + 1)的节点。

/**
 * Definition for a binary tree node.
 * public class TreeNode {
 *     int val;
 *     TreeNode left;
 *     TreeNode right;
 *     TreeNode(int x) { val = x; }
 * }
 */
class Solution {
    List<Integer> ans;
    TreeNode target;
    int K;
    public List<Integer> distanceK(TreeNode root, TreeNode target, int K) {
        ans = new LinkedList();
        this.target = target;
        this.K = K;
        dfs(root);
        return ans;
    }

    // Return vertex distance from node to target if exists, else -1
    // Vertex distance: the number of vertices on the path from node to target
    public int dfs(TreeNode node) {
        if (node == null)
            return -1;
        else if (node == target) {
            subtree_add(node, 0);
            return 1;
        } else {
            int L = dfs(node.left), R = dfs(node.right);
            if (L != -1) {
                if (L == K) ans.add(node.val);
                subtree_add(node.right, L + 1);
                return L + 1;
            } else if (R != -1) {
                if (R == K) ans.add(node.val);
                subtree_add(node.left, R + 1);
                return R + 1;
            } else {
                return -1;
            }
        }
    }

    // Add all nodes 'K - dist' from the node to answer.
    public void subtree_add(TreeNode node, int dist) {
        if (node == null) return;
        if (dist == K)
            ans.add(node.val);
        else {
            subtree_add(node.left, dist + 1);
            subtree_add(node.right, dist + 1);
        }
    }
}