[leetcode]863. All Nodes Distance K in Binary Tree
阿新 • • 發佈:2018-12-24
[leetcode]863. All Nodes Distance K in Binary Tree
Analysis
normal day—— [每天刷題並不難0.0]
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.
用map吧每個節點的父親節點記錄下來,方便往上搜索,然後再用DFS解決~
Implement
/**
* 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 res;
findParent(root);
dfs(target, K);
return res;
}
void findParent(TreeNode* node){
if(!node)
return;
if(node->left){
parent[node->left] = node;
findParent(node->left);
}
if(node->right){
parent[node->right] = node;
findParent(node->right);
}
}
void dfs(TreeNode* node, int k){
if(visit.find(node) != visit.end())
return ;
visit.insert(node);
if(k == 0)
res.push_back(node->val);
if(node->left)
dfs(node->left, k-1);
if(node->right)
dfs(node->right, k-1);
TreeNode* p = parent[node];
if(p)
dfs(p, k-1);
}
private:
vector<int> res;
map<TreeNode*, TreeNode*> parent;
set<TreeNode*> visit;
};