forked from kamyu104/LeetCode-Solutions
-
Notifications
You must be signed in to change notification settings - Fork 0
/
find-distance-in-a-binary-tree.cpp
87 lines (83 loc) · 2.73 KB
/
find-distance-in-a-binary-tree.cpp
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
// Time: O(n)
// Space: O(h)
/**
* Definition for a binary tree node.
* struct TreeNode {
* int val;
* TreeNode *left;
* TreeNode *right;
* TreeNode() : val(0), left(nullptr), right(nullptr) {}
* TreeNode(int x) : val(x), left(nullptr), right(nullptr) {}
* TreeNode(int x, TreeNode *left, TreeNode *right) : val(x), left(left), right(right) {}
* };
*/
class Solution {
public:
int findDistance(TreeNode* root, int p, int q) {
return iter_dfs(root, p, q);
}
private:
int iter_dfs(TreeNode *node, int p, int q) {
int result = 0, dist = -1;
vector<tuple<int, TreeNode *, unique_ptr<int>, unique_ptr<int>, int*>> stk;
stk.emplace_back(1, node, nullptr, nullptr, &dist);
while (!empty(stk)) {
const auto [step, node, ret1, ret2, ret] = move(stk.back()); stk.pop_back();
if (step == 1) {
if (!node) {
continue;
}
auto ret1 = make_unique<int>(-1), ret2 = make_unique<int>(-1);
auto p1 = ret1.get(), p2 = ret2.get();
stk.emplace_back(2, node, move(ret1), move(ret2), ret);
stk.emplace_back(1, node->right, nullptr, nullptr, p2);
stk.emplace_back(1, node->left, nullptr, nullptr, p1);
} else if (step == 2) {
if (node->val == p || node->val == q) {
if (*ret1 == -1 && *ret2 == -1) {
*ret = 0;
}
result = (*ret1 != -1) ? *ret1 + 1 : *ret2 + 1;
} else if (*ret1 != -1 && *ret2 != -1) {
result = *ret1 + *ret2 + 2;
} else if (*ret1 != -1) {
*ret = *ret1 + 1;
} else if (*ret2 != -1) {
*ret = *ret2 + 1;
}
}
}
return result;
}
};
// Time: O(n)
// Space: O(h)
class Solution2 {
public:
int findDistance(TreeNode* root, int p, int q) {
int result = 0;
dfs(root, p, q, &result);
return result;
}
private:
int dfs(TreeNode *node, int p, int q, int *result) {
if (!node) {
return -1;
}
auto left = dfs(node->left, p, q, result);
auto right = dfs(node->right, p, q, result);
if (node->val == p || node->val == q) {
if (left == -1 && right == -1) {
return 0;
}
*result = (left != -1) ? left + 1 : right + 1;
} else if (left != -1 && right != -1) {
*result = left + right + 2;
} else if (left != -1) {
return left + 1;
} else if (right != -1) {
return right + 1;
}
return -1;
}
};