#include <bits/stdc++.h>
using namespace std;
using ll = long long ;
#define pb emplace_back
#define AI(i) begin(i), end(i)
template<class T> bool chmin(T &a, T b) { return b < a && (a = b, true); }
template<class T> bool chmax(T &a, T b) { return a < b && (a = b, true); }
#ifdef KEV
#define DE(args...) kout("[ " + string(#args) + " ] = ", args)
void kout() { cerr << endl; }
template<class T, class ...U> void kout(T a, U ...b) { cerr << a << ' ', kout(b...); }
template<class T> void debug(T l, T r) { while (l != r) cerr << *l << " \n"[next(l) == r], ++l; }
#else
#define DE(...) 0
#define debug(...) 0
#endif
const int MAX_N = 500010, MAX_K = 20;
int n, k, belong[MAX_N];
vector<int> edge[MAX_N];
vector<int> state[MAX_N];
ll sum[MAX_N];
int anc[MAX_K][MAX_N], in[MAX_N], out[MAX_N];
bool isanc(int a, int b) { return in[a] <= in[b] && out[a] >= out[b]; }
void dfs(int x, int lst = 1) {
static int t;
anc[0][x] = lst;
in[x] = ++t;
for (int u : edge[x]) if (u != lst)
dfs(u, x);
out[x] = t;
}
void dfs2(int x, int lst = -1) {
for (int u : edge[x]) if (u != lst) {
dfs2(u, x);
sum[x] += sum[u];
}
}
int dfs3(int x, int lst = -1) {
int son = 0;
for (int u : edge[x]) if (u != lst)
son += dfs3(u, x);
if (son == 0 && sum[x] == 0 && lst != -1) {
++son;
DE(x);
}
return son;
}
int getlca(int a, int b) {
for (int d = MAX_K-1;d >= 0;--d)
if (!isanc(anc[d][a], b))
a = anc[d][a];
return isanc(a, b) ? a : anc[0][a];
}
int32_t main() {
ios_base::sync_with_stdio(0), cin.tie(0);
cin >> n >> k;
for (int a, b, i = 1;i < n;++i) {
cin >> a >> b;
edge[a].pb(b), edge[b].pb(a);
}
for (int i = 1;i <= n;++i) {
cin >> belong[i];
state[ belong[i] ].pb(i);
}
int rt = -1;
for (int i = 1;i <= k;++i) {
if (state[i].size() > 1)
rt = state[i][0];
}
if (rt == -1) {
return cout << k-1 << '\n', 0;
}
DE(rt);
dfs(rt);
for (int d = 1;d < MAX_K;++d)
for (int i = 1;i <= n;++i)
anc[d][i] = anc[d-1][ anc[d-1][i] ];
for (int i = 1;i <= k;++i) {
if (state[i].size() > 1) {
sort(AI(state[i]), [&](int a, int b) {
return in[a] < in[b];
});
for (int x : state[i])
++sum[x];
sum[ getlca(state[i][0], state[i].back()) ] -= state[i].size();
}
}
dfs2(rt);
int res = dfs3(rt);
if (res > 1) --res;
cout << res << '\n';
}
Compilation message
mergers.cpp: In function 'int dfs3(int, int)':
mergers.cpp:14:17: warning: statement has no effect [-Wunused-value]
14 | #define DE(...) 0
| ^
mergers.cpp:46:3: note: in expansion of macro 'DE'
46 | DE(x);
| ^~
mergers.cpp: In function 'int32_t main()':
mergers.cpp:14:17: warning: statement has no effect [-Wunused-value]
14 | #define DE(...) 0
| ^
mergers.cpp:76:2: note: in expansion of macro 'DE'
76 | DE(rt);
| ^~
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
16 ms |
24044 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
16 ms |
24044 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
16 ms |
24044 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Correct |
83 ms |
39268 KB |
Output is correct |
| 2 |
Incorrect |
82 ms |
32612 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
Verdict |
Execution time |
Memory |
Grader output |
| 1 |
Incorrect |
16 ms |
24044 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
