给定一个二叉树(具有根结点 root
), 一个目标结点 target
,和一个整数值 k
。
返回到目标结点 target
距离为 k
的所有结点的值的列表。 答案可以以 任何顺序 返回。
示例 1:
输入:root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
输出:[7,4,1]
解释:所求结点为与目标结点(值为 5)距离为 2 的结点,值分别为 7,4,以及 1
示例 2:
输入: root = [1], target = 1, k = 3
输出: []
深度优先搜索 + 哈希表
若将 target\textit{target}target 当作树的根结点,我们就能从 target\textit{target}target 出发,使用深度优先搜索去寻找与 target\textit{target}target 距离为 kkk 的所有结点,即深度为 kkk 的所有结点。
由于输入的二叉树没有记录父结点,为此,我们从根结点 root\textit{root}root 出发,使用深度优先搜索遍历整棵树,同时用一个哈希表记录每个结点的父结点。
然后从 target\textit{target}target 出发,使用深度优先搜索遍历整棵树,除了搜索左右儿子外,还可以顺着父结点向上搜索。
代码实现时,由于每个结点值都是唯一的,哈希表的键可以用结点值代替。此外,为避免在深度优先搜索时重复访问结点,递归时额外传入来源结点 from\textit{from}from,在递归前比较目标结点是否与来源结点相同,不同的情况下才进行递归。
class Solution {
unordered_map<int, TreeNode*> parents;
vector<int> ans;
void findParents(TreeNode* node) {
if (node->left != nullptr) {
parents[node->left->val] = node;
findParents(node->left);
}
if (node->right != nullptr) {
parents[node->right->val] = node;
findParents(node->right);
}
}
void findAns(TreeNode* node, TreeNode* from, int depth, int k) {
if (node == nullptr) {
return;
}
if (depth == k) {
ans.push_back(node->val);
return;
}
if (node->left != from) {
findAns(node->left, node, depth + 1, k);
}
if (node->right != from) {
findAns(node->right, node, depth + 1, k);
}
if (parents[node->val] != from) {
findAns(parents[node->val], node, depth + 1, k);
}
}
public:
vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
// 从 root 出发 DFS,记录每个结点的父结点
findParents(root);
// 从 target 出发 DFS,寻找所有深度为 k 的结点
findAns(target, nullptr, 0, k);
return ans;
}
};