题目链接:
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);
}
};