863. All Nodes Distance K in Binary Tree
https://leetcode.com/problems/all-nodes-distance-k-in-binary-tree/
/*
* @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
Last updated