leetcode863. 二叉树中所有距离为 K 的结点

题目链接：
https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree/
题目：
给定一个二叉树（具有根结点 root）， 一个目标结点 target ，和一个整数值 K 。

返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。

示例 1：

输入：root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2

输出：[7,4,1]

解释：
所求结点为与目标结点（值为 5）距离为 2 的结点，
值分别为 7，4，以及 1

思路：
核心思想：利用map标记从root到target的路径上每个节点距离target的长度，如在本题中mp[target]=0,mp[root]=1，路径上面没有其他节点。然后从根节点开始搜索，遇到之前路径上的节点，就修改当前已经记录的路径长度。
code：

class Solution {
private:
    unordered_map<TreeNode*,int> mp;
public:
    vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
        vector<int> ans;
        if(root==NULL) return ans;
        findPath(root, target);
        dfs(root,target,K,mp[root],ans);
        return ans;
    }
    int findPath(TreeNode* root, TreeNode* target){//利用map存储从root到target的路径长度
        if(root==NULL) return -1;
        if(root==target){
            mp[root]=0;
            return 0;
        }
        int left=findPath(root->left,target);
        if(left>=0){
            mp[root]=left+1;
            return left+1;
        }
        int right=findPath(root->right,target);
        if(right>=0){
            mp[root]=right+1;
            return right+1;
        }
        
        return -1;
    }
    
     void dfs(TreeNode *root,TreeNode* target, int K,int length,vector<int> &ans){
        if(root==NULL) return;
        if(mp.count(root)) length=mp[root];//如果在root到target的路径上，就更换为原来的路径，注意这里不能用mp[root]==0来判断，一来这种操作本身很危险，其次因为mp[target]=0会把target的情况忽略掉...
        
        if(K==length) ans.push_back(root->val);
        
        dfs(root->left,target,K,length+1,ans);
        dfs(root->right,target,K,length+1,ans);
    }
};