863. All Nodes Distance K in Binary Tree

Previous852. Peak Index in a Mountain ArrayNext889. Construct Binary Tree from Preorder and Postorder Traversal
Last updated

Last updated
Input: root = [3,5,1,6,2,0,8,null,null,7,4], target = 5, K = 2
Output: [7,4,1]
Explanation:
The nodes that are a distance 2 from the target node (with value 5)
have values 7, 4, and 1.
Note that the inputs "root" and "target" are actually TreeNodes.
The descriptions of the inputs above are just serializations of these objects.// DFS + Hashmap + Hashset
void buildParent(TreeNode* node, unordered_map<TreeNode*, TreeNode*>& parent) {
if (!node) return;
if (node->left) parent[node->left] = node;
if (node->right) parent[node->right] = node;
buildParent(node->left, parent);
buildParent(node->right, parent);
}
void helper(TreeNode* node, int K, unordered_map<TreeNode*, TreeNode*>& parent, unordered_set<TreeNode*>& visited, vector<int>& res) {
if (visited.count(node)) return;
visited.insert(node);
if (K == 0) {
res.push_back(node->val);
return;
}
if (node->left) helper(node->left, K - 1, parent, visited, res); // leftward
if (node->right) helper(node->right, K - 1, parent, visited, res); // rightward
if (parent[node]) helper(parent[node], K - 1, parent, visited, res); // backward
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int K) { // time: O(n); space: O(n)
unordered_map<TreeNode*, TreeNode*> parent;
unordered_set<TreeNode*> visited;
vector<int> res;
buildParent(root, parent);
helper(target, K, parent, visited, res);
return res;
}// BFS Level Order Traversal
void buildMap(TreeNode* node, TreeNode* pre, unordered_map<TreeNode*, vector<TreeNode*> >& mp) {
if (!node || mp.count(node)) return;
if (pre) {
mp[node].push_back(pre);
mp[pre].push_back(node);
}
buildMap(node->left, node, mp);
buildMap(node->right, node, mp);
}
vector<int> distanceK(TreeNode* root, TreeNode* target, int K) { // time: O(n); space: O(n)
if (!root) return {};
vector<int> res;
unordered_map<TreeNode*, vector<TreeNode*> > mp;
queue<TreeNode*> q({target});
unordered_set<TreeNode*> visited({target});
buildMap(root, nullptr, mp);
while (!q.empty()) {
if (K == 0) {
for (int i = q.size() - 1; i >= 0; --i) {
res.push_back(q.front()->val); q.pop();
}
break;
}
for (int i = q.size() - 1; i >= 0; --i) {
TreeNode* t = q.front(); q.pop();
for (TreeNode*& node : mp[t]) {
if (visited.count(node)) continue;
q.push(node);
visited.insert(node);
}
}
--K;
}
return res;
}