863. 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:
map<TreeNode*, TreeNode*> parent;
void dfs(TreeNode* root) {
if(!root) return ;
if(root->left) parent[root->left] = root;
if(root->right) parent[root->right] = root;
dfs(root->left);
dfs(root->right);
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
dfs(root);
parent[root] = NULL;
vector<int> V;
int visited[1010] = {0};
queue<TreeNode*> Q;
TreeNode* p;
if(target) Q.push(target);
int cnt = 0;
while(!Q.empty() && cnt <= K) {
int len = Q.size();
for(int ii = 0; ii < len; ii++) {
p = Q.front();
Q.pop();
if(!visited[p->val]) {
if(cnt == K) V.push_back(p->val);
if(p->left) Q.push(p->left);
if(p->right) Q.push(p->right);
if(parent[p]) Q.push(parent[p]);
visited[p->val] = 1;
}
}
cnt++;
}
return V;
}
};