#include <bits/stdc++.h>
using namespace std;
#define ll long long
#define N (ll)(5e5+5)
#define fi first
#define se second
ll n, k, sk[N], p[21][N], st[N], ft[N], t[N], ct, gp[N], d[N], ans, part[N], sz[N], f;
vector<ll> g[N], u[N];
void dfs(ll s, ll par) {
st[s] = ++ct; u[sk[s]].push_back(s); p[0][s] = par; for (int i = 1; i < 21; i++) p[i][s] = p[i - 1][p[i - 1][s]];
for (auto x : g[s]) if (x != par) dfs(x, s); ft[s] = ++ct;
}
inline ll glca(ll x, ll y) {
if (x == y) return x;
if (st[x] < st[y]) swap(x, y);
for (int i = 20; i > -1; i--) if (!(st[p[i][x]] <= st[y] && ft[p[i][x]] >= ft[y]))x = p[i][x];
return p[0][x];
}
void dfs2(ll s, ll par) { for (auto x : g[s]) if (x != par) dfs2(x, s), t[s] += t[x]; }
ll gpar(ll x) { return part[x] == part[part[x]] ? part[x] : part[x] = gpar(part[x]); }inline void merge(ll x, ll y) { x = gpar(x), y = gpar(y); if (x == y) return; if (sz[x] < sz[y]) swap(x, y); part[y] = x, sz[x] += sz[y]; f--; }
//void dfss(ll s, ll par) { for (auto x : g[s]) if (x != par) { if (t[x] == 0) gp[x] = ++ct, d[gp[s]]++, d[ct]++; else gp[x] = gp[s]; dfss(x, s); } }
void cnt(ll s,ll par){
for (auto x : g[s]) if (x != par) cnt(x, s);
par = gpar(par); s = gpar(s); if (par != s) { d[s]++, d[par]++; if (d[s] == 2) ans--; if (d[par] == 2) ans--; if (d[s] == 1) ans++; if (d[par] == 1) ans++;}
}
int main() {
ios_base::sync_with_stdio(false), cin.tie(0), cout.tie(0);
ll a, b; cin >> n >> k; for (int i = 1; i < n; i++) cin >> a >> b, g[a].push_back(b), g[b].push_back(a);
for (int i = 1; i <= n; i++) cin >> sk[i], part[i] = i, sz[i] = 1; if (n == 1) return cout << 0, 0;
dfs(1, 1);
for (int i = 1; i <= k; i++) for (int j = 1; j < u[i].size(); j++) ++t[u[i][j]], ++t[u[i][j - 1ll]], ----t[glca(u[i][j], u[i][j - 1ll])];
f = n; ans = 0; ct = 1; gp[1] = 1; dfs2(1, 1); for (int i = 1; i <= n; i++) if (t[i] != 0) merge(i, p[0][i]);
if (f == 1) return cout << 0, 0; cnt(1, 1); cout << ans - 1;
}
Compilation message
mergers.cpp: In function 'void dfs(long long int, long long int)':
mergers.cpp:11:2: warning: this 'for' clause does not guard... [-Wmisleading-indentation]
11 | for (auto x : g[s]) if (x != par) dfs(x, s); ft[s] = ++ct;
| ^~~
mergers.cpp:11:47: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'for'
11 | for (auto x : g[s]) if (x != par) dfs(x, s); ft[s] = ++ct;
| ^~
mergers.cpp: In function 'int main()':
mergers.cpp:29:2: warning: this 'for' clause does not guard... [-Wmisleading-indentation]
29 | for (int i = 1; i <= n; i++) cin >> sk[i], part[i] = i, sz[i] = 1; if (n == 1) return cout << 0, 0;
| ^~~
mergers.cpp:29:69: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'for'
29 | for (int i = 1; i <= n; i++) cin >> sk[i], part[i] = i, sz[i] = 1; if (n == 1) return cout << 0, 0;
| ^~
mergers.cpp:31:49: warning: comparison of integer expressions of different signedness: 'int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
31 | for (int i = 1; i <= k; i++) for (int j = 1; j < u[i].size(); j++) ++t[u[i][j]], ++t[u[i][j - 1ll]], ----t[glca(u[i][j], u[i][j - 1ll])];
| ~~^~~~~~~~~~~~~
mergers.cpp:33:2: warning: this 'if' clause does not guard... [-Wmisleading-indentation]
33 | if (f == 1) return cout << 0, 0; cnt(1, 1); cout << ans - 1;
| ^~
mergers.cpp:33:35: note: ...this statement, but the latter is misleadingly indented as if it were guarded by the 'if'
33 | if (f == 1) return cout << 0, 0; cnt(1, 1); cout << ans - 1;
| ^~~
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
13 ms |
24024 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
13 ms |
24024 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
13 ms |
24024 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Correct |
243 ms |
50132 KB |
Output is correct |
| 2 |
Incorrect |
126 ms |
52224 KB |
Output isn't correct |
| 3 |
Halted |
0 ms |
0 KB |
- |
| # |
결과 |
실행 시간 |
메모리 |
Grader output |
| 1 |
Incorrect |
13 ms |
24024 KB |
Output isn't correct |
| 2 |
Halted |
0 ms |
0 KB |
- |
