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.
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; } };
|