给定一个二叉树(具有根结点 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注意,输入的 "root" 和 "target" 实际上是树上的结点。
上面的输入仅仅是对这些对象进行了序列化描述。提示:
给定的树是非空的,且最多有 K 个结点。
树上的每个结点都具有唯一的值 0 <= node.val <= 500 。
目标结点 target 是树上的结点。
0 <= K <= 1000.来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/all-nodes-distance-k-in-binary-tree
著作权归领扣网络所有。商业转载请联系官方授权,非商业转载请注明出处。
解法:
/**
* 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 {};
unordered_map<TreeNode*, vector<TreeNode *>> mp;
queue<TreeNode *> qu;
solution(root, nullptr, mp);
qu.push(target);
set<TreeNode *> s;
s.insert(target);
vector<int> res;
while(!qu.empty())
{
if(K == 0)
{
while(!qu.empty())
{
res.push_back(qu.front()->val);
qu.pop();
}
return res;
}
int n = qu.size();
while(n--)
{
auto it = qu.front();
qu.pop();
for(auto c : mp[it])
{
if(s.find(c) != s.end())
continue;
qu.push(c);
s.insert(c);
}
}
K--;
}
return {};
}
void solution(TreeNode *root, TreeNode *pre, unordered_map<TreeNode*, vector<TreeNode *>> &mp)
{
if(root == nullptr) return;
if(pre && root)
{
mp[pre].push_back(root);
mp[root].push_back(pre);
}
solution(root->left, root, mp);
solution(root->right, root, mp);
}
};