Leetcode 863. All Nodes Distance K in Binary Tree
Given the root of a binary tree, the value of a target node in the tree, and an integer k
, return an array of the values of all nodes that have a distance k
from the target node.
You can return the answer in any order.
Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, k = 2
Output: [7,4,1]
k
?
val
, left
, and right
.k
distance away.k
.Here is the implementation in C++:
#include <vector>
#include <unordered_map>
#include <unordered_set>
#include <queue>
using namespace std;
// Definition for a binary tree node.
struct TreeNode {
int val;
TreeNode *left;
TreeNode *right;
TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
};
class Solution {
private:
unordered_map<int, vector<int>> graph;
void buildGraph(TreeNode* node) {
if (!node) return;
if (node->left) {
graph[node->val].push_back(node->left->val);
graph[node->left->val].push_back(node->val);
buildGraph(node->left);
}
if (node->right) {
graph[node->val].push_back(node->right->val);
graph[node->right->val].push_back(node->val);
buildGraph(node->right);
}
}
public:
vector<int> distanceK(TreeNode* root, TreeNode* target, int k) {
if (!root) return {};
buildGraph(root);
unordered_set<int> visited;
queue<int> q;
q.push(target->val);
visited.insert(target->val);
int distance = 0;
while (!q.empty()) {
int size = q.size();
if (distance == k) {
vector<int> result;
while (!q.empty()) {
result.push_back(q.front());
q.pop();
}
return result;
}
for (int i = 0; i < size; ++i) {
int node = q.front();
q.pop();
for (int neighbor : graph[node]) {
if (visited.find(neighbor) == visited.end()) {
visited.insert(neighbor);
q.push(neighbor);
}
}
}
++distance;
}
return {};
}
};
N
is the number of nodes in the tree. Each node and its children are visited once.Overall, the time complexity is O(N).
The space complexity is also O(N) due to the space used to store the graph and the BFS queue.
Got blindsided by a question you didn’t expect?
Spend too much time studying?
Or simply don’t have the time to go over all 3000 questions?