863. 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.
c++1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
|
class Solution {
struct Node {
int val;
list<Node*> near;
Node(int x) : val(x) {}
Node(int x, Node* n) : val(x), near{n} {}
};
Node* buildGraph(TreeNode* root, TreeNode* target) {
Node* t;
queue<pair<TreeNode*,Node*>> q;
q.push({root, new Node(root->val)});
while(!q.empty()) {
auto p = q.front();
q.pop();
if(p.first == target) t = p.second;
if(p.first->left != NULL) {
Node* left = new Node(p.first->left->val, p.second);
p.second->near.push_back(left);
q.push({p.first->left, left});
}
if(p.first->right != NULL) {
Node* right = new Node(p.first->right->val, p.second);
p.second->near.push_back(right);
q.push({p.first->right, right});
}
}
return t;
}
public:
vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
if(!k) return {target->val};
Node* t = buildGraph(root, target);
vector<int> res;
queue<Node*> q;
set<Node*> v;
q.push(t), v.insert(t);
while(k--) {
int sz = q.size();
while(sz--) {
auto n = q.front();
q.pop();
for(auto near : n->near) {
if(v.count(near)) continue;
v.insert(near);
q.push(near);
}
}
}
while(!q.empty()) {
res.push_back(q.front()->val);
q.pop();
}
return res;
}
};
|
