给定一个二叉树(具有根结点 root), 一个目标结点 target ,和一个整数值 K 。
返回到目标结点 target 距离为 K 的所有结点的值的列表。 答案可以以任何顺序返回。
来源:力扣(LeetCode)
链接:https://leetcode-cn.com/problems/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 {
private:
vector<int> ans;
unordered_map <TreeNode*, TreeNode*> mp;
public:
void findParents(TreeNode* node) {
if (node->left) {
mp[node->left] = node;
findParents(node->left);
}
if (node->right) {
mp[node->right] = node;
findParents(node->right);
}
}
void dfs(TreeNode* node, TreeNode* from, int depth, int k) {
if (!node) {
return;
}
if (depth == k) {
ans.push_back(node->val);
return;
}
// != from 表示单向传播
if (node->left != from) {
dfs(node->left, node, depth + 1, k);
}
if (node->right != from) {
dfs(node->right, node, depth + 1, k);
}
if (mp[node] != from) {
dfs(mp[node], node, depth + 1, k);
}
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
findParents(root);
dfs(target, nullptr, 0, k);
return ans;
}
};