#include<bits/stdc++.h>
#define mp make_pair
#define F first
#define S second
#define iter(a) a.begin(), a.end()
#define lsort(a) sort(iter(a))
#define printv(a, b) { \
for(auto pv : a) b << pv << " "; \
b << "\n"; \
}
#define eb emplace_back
using namespace std;
typedef long long ll;
using pll = pair<ll, ll>;
ostream& operator<<(ostream& o, pll p){
return o << '(' << p.F << ',' << p.S << ')';
}
void waassert(bool t){
if(!t){
cout << "OAO\n";
exit(0);
}
}
struct DSU{
vector<int> dsu;
int cnt;
void init(int n){
dsu.resize(n + 1);
cnt = n;
for(int i = 1; i <= n; i++) dsu[i] = i;
}
int findDSU(int a){
if(dsu[a] != a) dsu[a] = findDSU(dsu[a]);
return dsu[a];
}
void unionDSU(int a, int b){
a = findDSU(a);
b = findDSU(b);
if(a == b) return;
cnt--;
dsu[b] = a;
}
};
int n;
vector<vector<int>> g;
vector<int> s;
vector<vector<int>> anc;
vector<int> in, out;
DSU dsu;
int ts = 0;
void dfs1(int now, int p){
anc[0][now] = p;
in[now] = ts++;
for(int i : g[now]){
if(i == p) continue;
dfs1(i, now);
}
out[now] = ts++;
}
void buildLCA(){
for(int i = 1; i < 20; i++){
for(int j = 1; j <= n; j++){
anc[i][j] = anc[i - 1][anc[i - 1][j]];
}
}
}
bool isAnc(int a, int b){
return in[a] <= in[b] && out[a] >= out[b];
}
int getLCA(int a, int b){
if(isAnc(a, b)) return a;
for(int i = 19; i >= 0; i--){
if(!isAnc(anc[i][a], b)) a = anc[i][a];
}
return anc[a][0];
}
int unionPath(int a, int b){
int lca = getLCA(a, b);
while(!isAnc(a, lca)){
dsu.unionDSU(anc[0][a], a);
a = dsu.findDSU(a);
}
while(!isAnc(b, lca)){
dsu.unionDSU(anc[0][b], b);
b = dsu.findDSU(b);
}
return dsu.findDSU(lca);
}
int fv = -1, fd = -1;
void dfs2(int now, int p, int dpt){
if(dpt > fd) fv = now, fd = dpt;
//cerr << "dfs2 " << now << " " << p << " " << dpt << "\n";
for(int i : g[now]){
if(i == p) continue;
if(dsu.findDSU(now) != dsu.findDSU(i)) dfs2(i, now, dpt + 1);
else dfs2(i, now, dpt);
}
}
void owo(){
//cerr << "owo\n";
fv = -1; fd = -1;
dfs2(1, 1, 0);
int t1 = fv;
fv = -1; fd = -1;
dfs2(t1, t1, 0);
int t2 = fv;
//cerr << t1 << " " << t2 << "\n";
unionPath(t1, t2);
}
int main(){
ios_base::sync_with_stdio(false);
cin.tie(0);
int k;
cin >> n >> k;
g.resize(n + 1);
s.resize(n + 1);
anc.resize(20, vector<int>(n + 1));
in.resize(n + 1);
out.resize(n + 1);
dsu.init(n);
vector<vector<int>> city(k + 1);
for(int i = 0; i < n - 1; i++){
int u, v;
cin >> u >> v;
g[u].eb(v);
g[v].eb(u);
}
for(int i = 1; i <= n; i++){
cin >> s[i];
city[s[i]].eb(i);
}
dfs1(1, 1);
buildLCA();
for(int i = 1; i <= k; i++){
waassert(city[i].size() >= 1);
int v = city[i].front();
for(int j = 1; j < city[i].size(); j++){
v = unionPath(v, city[i][j]);
}
}
int ans = 0;
while(dsu.cnt > 1){
owo();
ans++;
}
cout << ans << "\n";
return 0;
}
Compilation message
mergers.cpp: In function 'int main()':
mergers.cpp:160:26: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
160 | for(int j = 1; j < city[i].size(); j++){
| ~~^~~~~~~~~~~~~~~~
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Runtime error |
1 ms |
492 KB |
Execution killed with signal 11 |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Runtime error |
1 ms |
492 KB |
Execution killed with signal 11 |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Runtime error |
1 ms |
492 KB |
Execution killed with signal 11 |
3 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Runtime error |
74 ms |
32740 KB |
Execution killed with signal 11 |
2 |
Halted |
0 ms |
0 KB |
- |
# |
결과 |
실행 시간 |
메모리 |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Runtime error |
1 ms |
492 KB |
Execution killed with signal 11 |
3 |
Halted |
0 ms |
0 KB |
- |