I'm posting two similar solutions for LeetCode's "All Nodes Distance K in Binary Tree". If you'd like to review, please do so. Thank you!
Problem
We are given a binary tree (with root node
root
), atarget
node, and an integer valueK
.Return a list of the values of all nodes that have a distance
K
from the target node. The answer can be returned in any order.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.
Note:
- The given tree is non-empty.
- Each node in the tree has unique values 0 <= node.val <= 500.
- The target node is a node in the tree.
- 0 <= K <= 1000.
Note that the inputs "root" and "target" are actually TreeNodes. The descriptions of the inputs above are just serializations of these objects.
Solution 1
// The following block might slightly improve the execution time;
// Can be removed;
static const auto __optimize__ = []() {
std::ios::sync_with_stdio(false);
std::cin.tie(nullptr);
std::cout.tie(nullptr);
return 0;
}();
// Most of headers are already included;
// Can be removed;
#include <cstdint>
#include <vector>
#include <unordered_map>
#include <unordered_set>
using ValueType = int;
static const struct Solution {
static const std::vector<ValueType> distanceK(
TreeNode* root,
TreeNode* target,
const ValueType K
) {
std::vector<ValueType> res;
std::unordered_map<TreeNode*, TreeNode*> parents;
std::unordered_set<TreeNode*> visited;
getParent(root, parents);
depthFirstSearch(target, K, parents, visited, res);
return res;
}
private:
static const void getParent(
TreeNode* node,
std::unordered_map<TreeNode*, TreeNode*>& parents
) {
if (!node) {
return;
}
if (node->left) {
parents[node->left] = node;
getParent(node->left, parents);
}
if (node->right) {
parents[node->right] = node;
getParent(node->right, parents);
}
}
static const void depthFirstSearch(
TreeNode* node,
const ValueType K,
std::unordered_map<TreeNode*, TreeNode*>& parents,
std::unordered_set<TreeNode*>& visited,
std::vector<ValueType>& res
) {
if (!node) {
return;
}
if (visited.count(node) > 0) {
return;
}
visited.insert(node);
if (!K) {
res.emplace_back(node->val);
return;
}
depthFirstSearch(node->left, K - 1, parents, visited, res);
depthFirstSearch(node->right, K - 1, parents, visited, res);
depthFirstSearch(parents[node], K - 1, parents, visited, res);
}
};
Solution 2
// The following block might slightly improve the execution time;
// Can be removed;
static const auto __optimize__ = []() {
std::ios::sync_with_stdio(false);
std::cin.tie(nullptr);
std::cout.tie(nullptr);
return 0;
}();
// Most of headers are already included;
// Can be removed;
#include <cstdint>
#include <vector>
#include <unordered_map>
#include <unordered_set>
using ValueType = int;
static const struct Solution {
const std::vector<ValueType> distanceK(
TreeNode* root,
TreeNode* target,
ValueType K
) {
getParent(root);
depthFirstSearch(target, K);
return res;
}
private:
std::vector<ValueType> res;
std::unordered_map<TreeNode*, TreeNode*> parents;
std::unordered_set<TreeNode*> visited;
const void getParent(
TreeNode* node
) {
if (!node) {
return;
}
if (node->left) {
parents[node->left] = node;
getParent(node->left);
}
if (node->right) {
parents[node->right] = node;
getParent(node->right);
}
}
const void depthFirstSearch(
TreeNode* node,
const ValueType K
) {
if (!node) {
return;
}
if (visited.count(node) > 0) {
return;
}
visited.insert(node);
if (!K) {
res.emplace_back(node->val);
return;
}
depthFirstSearch(node->left, K - 1);
depthFirstSearch(node->right, K - 1);
depthFirstSearch(parents[node], K - 1);
}
};
Reference
Here is LeetCode's template:
/**
* 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> distanceK(TreeNode* root, TreeNode* target, int K) {
}
};