863. All Nodes Distance K in Binary Tree
阿新 • • 發佈:2018-12-04
-- tar false pop while order roo back target
/** * 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) { vector<bool> visited(501, false); vector<unordered_set<int>> g(501); dfs(root, g); vector<int> res; queue<int> q; q.push(target->val); visited[target->val] = true; int lv = 0; while (!q.empty()) {int qs = q.size(); if (lv == K) { while(!q.empty()) { res.push_back(q.front()); q.pop(); } break; } while (qs-- > 0) { int p = q.front(); q.pop();for (auto i : g[p]) { if (!visited[i]) { q.push(i); visited[i] = true; } } } lv++; } return res; } void dfs(TreeNode *root, vector<unordered_set<int>>& g) { if (root == NULL) return; if (root->left) { g[root->val].insert(root->left->val); g[root->left->val].insert(root->val); } if (root->right) { g[root->val].insert(root->right->val); g[root->right->val].insert(root->val); } dfs(root->left, g); dfs(root->right, g); } };
863. All Nodes Distance K in Binary Tree