-
Notifications
You must be signed in to change notification settings - Fork 111
/
minimum-height-trees.cc
78 lines (74 loc) · 1.74 KB
/
minimum-height-trees.cc
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
// Minimum Height Trees
#define REP(i, n) for (int i = 0; i < (n); i++)
class Solution {
public:
vector<int> findMinHeightTrees(int n, const vector<vector<int>>& edges) {
vector<vector<int>> es(n);
vector<int> d(n, 0);
vector<bool> removed(n, false);
for (auto &e: edges) {
es[e[0]].push_back(e[1]);
es[e[1]].push_back(e[0]);
d[e[0]]++;
d[e[1]]++;
}
vector<int> leaves;
REP(i, n)
if (d[i] <= 1)
leaves.push_back(i);
while (n > 2) {
vector<int> leaves2;
n -= leaves.size();
for (int u: leaves) {
int p = -1;
for (int v: es[u])
if (! removed[v])
p = v;
removed[u] = true;
if (--d[p] == 1)
leaves2.push_back(p);
}
leaves.swap(leaves2);
}
return leaves;
}
};
/// find a diameter
class Solution {
vector<vector<int>> es;
vector<int> parent;
int bfs(int u) {
queue<int> q;
int last = u;
fill(parent.begin(), parent.end(), -1);
parent[u] = -2;
q.push(u);
while (! q.empty()) {
u = q.front();
q.pop();
last = u;
for (int v: es[u])
if (parent[v] == -1) {
parent[v] = u;
q.push(v);
}
}
return last;
}
public:
vector<int> findMinHeightTrees(int n, vector<pair<int, int>>& edges) {
es.resize(n);
parent.resize(n);
for (auto &e: edges) {
es[e.first].push_back(e.second);
es[e.second].push_back(e.first);
}
vector<int> ancestors;
for (int v = bfs(bfs(0)); v != -2; v = parent[v])
ancestors.push_back(v);
vector<int> ret{ancestors[ancestors.size()/2]};
if (ancestors.size() % 2 == 0)
ret.push_back(ancestors[ancestors.size()/2-1]);
return ret;
}
};