863. All Nodes Distance K in Binary Tree

题目：

Given the root of a binary tree, the value of a target node target, and an integer k, return an array of the values of all nodes that have a distance k from the target node.

You can return the answer 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.

Example 2:

Input: root = [1], target = 1, k = 3
Output: []

Constraints:

• The number of nodes in the tree is in the range [1, 500].
• 0 <= Node.val <= 500
• All the values Node.val are unique.
• target is the value of one of the nodes in the tree.
• 0 <= k <= 1000

思路：

这道题一看到肯定是想把树掰直，比如例子，想把它变成多个数组(比如[7, 2, 5, 3, 1, 8], [6, 5, 3, 1, 0] 等等)然后一一寻找5周围2个index的数。但是其实这样会有很多重复查找，即使实现了复杂度也比较高。但是由这个启发，如果我们是从target，即5开始寻找，往三个方向，分别是6，2，3来寻找，不是就能达到一样的效果了吗？为了避免重建树，用一个哈希map <int, TreeNode*>来记录每个结点的父节点，那么相当于每个结点最多有三个“连接点”，左右和父。之后只要把target当作root，用dfs查找depth为k的点即可。本题中记录父节点用的是BFS，寻找depth用的是DFS。

 

 

代码：

/**
 * Definition for a binary tree node.
 * struct TreeNode {
 *     int val;
 *     TreeNode *left;
 *     TreeNode *right;
 *     TreeNode(int x) : val(x), left(NULL), right(NULL) {}
 * };
 */
class Solution {
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
        if (!root)
            return {};
        queue<TreeNode*> q;
        q.push(root);
        while(q.size()) {
            int len =q.size();
            for(int i = 0; i < len; i++) {
                auto t = q.front();
                q.pop();
                if(t->left) {
                    m[t->left->val] = t;
                    q.push(t->left);
                }
                if(t->right) {
                    m[t->right->val] = t;
                    q.push(t->right);
                }
            }
        }
        depth = k;
        dfs(target, nullptr, 0);
        return res;
    }
private:
    unordered_map<int, TreeNode*> m;
    vector<int> res;
    int depth;
    void dfs(TreeNode* root, TreeNode* prev, int k) {
        if (k > depth)
            return;
        if (k == depth) {
            res.push_back(root->val);
            return;
        }
        if (root->left && root->left != prev) 
            dfs(root->left, root, k + 1);
        if (root->right && root->right != prev) 
            dfs(root->right, root, k + 1);
        if (m.count(root->val) && m[root->val] != prev) 
            dfs(m[root->val], root, k + 1);
    }
};