Skip to content

Commit 538e756

Browse files
authored
Create diameter-of-n-ary-tree.cpp
1 parent 8ce7ba2 commit 538e756

File tree

1 file changed

+72
-0
lines changed

1 file changed

+72
-0
lines changed

C++/diameter-of-n-ary-tree.cpp

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
// Time: O(n)
2+
// Space: O(h)
3+
4+
/*
5+
// Definition for a Node.
6+
class Node {
7+
public:
8+
int val;
9+
vector<Node*> children;
10+
11+
Node() {}
12+
13+
Node(int _val) {
14+
val = _val;
15+
}
16+
17+
Node(int _val, vector<Node*> _children) {
18+
val = _val;
19+
children = _children;
20+
}
21+
};
22+
*/
23+
24+
// Time: O(n)
25+
// Space: O(h)
26+
class Solution2 {
27+
public:
28+
int diameter(Node* root) {
29+
return iter_dfs(root).first;
30+
}
31+
32+
private:
33+
pair<int, int> iter_dfs(Node *root) {
34+
using RET = pair<int, int>;
35+
RET result;
36+
vector<tuple<int, Node *, shared_ptr<RET>, RET *>> stk = {{1, root, nullptr, &result}};
37+
while (!stk.empty()) {
38+
const auto [step, node, ret2, ret] = stk.back(); stk.pop_back();
39+
if (step == 1) {
40+
for (int i = node->children.size() - 1; i >= 0; --i) {
41+
const auto& ret2 = make_shared<RET>();
42+
stk.emplace_back(2, nullptr, ret2, ret);
43+
stk.emplace_back(1, node->children[i], nullptr, ret2.get());
44+
}
45+
} else {
46+
ret->first = max(ret->first, max(ret2->first, ret->second + ret2->second + 1));
47+
ret->second = max(ret->second, ret2->second + 1);
48+
}
49+
}
50+
return result;
51+
}
52+
};
53+
54+
// Time: O(n)
55+
// Space: O(h)
56+
class Solution2 {
57+
public:
58+
int diameter(Node* root) {
59+
return dfs(root).first;
60+
}
61+
62+
private:
63+
pair<int, int> dfs(Node *node) {
64+
int max_dia = 0, max_depth = 0;
65+
for (const auto& child : node->children) {
66+
const auto& [child_max_dia, child_max_depth] = dfs(child);
67+
max_dia = max({max_dia, child_max_dia, max_depth + child_max_depth + 1});
68+
max_depth = max(max_depth, child_max_depth + 1);
69+
}
70+
return {max_dia, max_depth};
71+
}
72+
};

0 commit comments

Comments
 (0)