题目描述
用一个哈希表记录 节点->父节点的映射
再用一个哈希表,记录状态。
由于已经有了第一个哈希表,直接从target节点开始搜索, 搜索三个方向,俩儿子和父亲。
C++ 代码
class Solution {
public:
vector<int> ans;
unordered_map<TreeNode*, TreeNode*> parent; // son=>parent
unordered_map<TreeNode*, int> visit; //record visied node
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);
}
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
if (!root) return {}; // 原函数
findParent(root); // 记录所有son和parent
dfs(target, K); // 从target进行
return ans;
}
void dfs(TreeNode* node, int K)
{ // 唯一逻辑
if(visit[node])
return;
visit[node] ++;
if(K == 0)
{
ans.push_back(node->val);
return;
}
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);
}
};