二叉树中找与给定结点相距K的所有结点。二叉树是一种特殊的图,因此可以把二叉树存成无向图然后用BFS找所有相距K的点。还有可以利用二叉树的性质做DFS,对target的每个ancestor(与t距离为d)收集它在另一个子树中与ancestor距离为K-d-1的所有结点。
/*
* @lc app=leetcode id=863 lang=cpp
*
* [863] All Nodes Distance K in Binary Tree
*/
// @lc code=start
/**
* 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> res;
void collect(TreeNode *cur, int d) {
if (cur == nullptr) return;
if (d == 0) {
res.push_back(cur->val);
return;
}
collect(cur->left, d - 1);
collect(cur->right, d - 1);
}
int dist(TreeNode *cur, TreeNode *target, const int K) {
if (cur == nullptr) return -1;
if (cur == target) {
collect(cur, K);
return 0;
}
int l = dist(cur->left, target, K);
if (l != -1) {
++l;
if (K - l > 0) {
collect(cur->right, K - l - 1);
} else if (K == l) {
res.push_back(cur->val);
}
return l;
}
int r = dist(cur->right, target, K);
if (r != -1) {
++r;
if (K - r > 0) {
collect(cur->left, K - r - 1);
} else if (K == r) {
res.push_back(cur->val);
}
return r;
}
return -1;
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int K) {
dist(root, target, K);
return res;
}
};
// @lc code=end