/*
ID: varunra2
LANG: C++
TASK: mergers
*/
#include <bits/stdc++.h>
using namespace std;
#ifdef DEBUG
#include "lib/debug.h"
#define debug(...) cerr << "[" << #__VA_ARGS__ << "]:", debug_out(__VA_ARGS__)
#define debug_arr(...) \
cerr << "[" << #__VA_ARGS__ << "]:", debug_arr(__VA_ARGS__)
#pragma GCC diagnostic ignored "-Wsign-compare"
//#pragma GCC diagnostic ignored "-Wunused-parameter"
//#pragma GCC diagnostic ignored "-Wunused-variable"
#else
#define debug(...) 42
#endif
#define EPS 1e-9
#define IN(A, B, C) assert(B <= A && A <= C)
#define INF (int)1e9
#define MEM(a, b) memset(a, (b), sizeof(a))
#define MOD 1000000007
#define MP make_pair
#define PB push_back
#define all(cont) cont.begin(), cont.end()
#define rall(cont) cont.end(), cont.begin()
#define x first
#define y second
const double PI = acos(-1.0);
typedef long long ll;
typedef long double ld;
typedef pair<int, int> PII;
typedef map<int, int> MPII;
typedef multiset<int> MSETI;
typedef set<int> SETI;
typedef set<string> SETS;
typedef vector<int> VI;
typedef vector<PII> VII;
typedef vector<VI> VVI;
typedef vector<string> VS;
#define rep(i, a, b) for (int i = a; i < (b); ++i)
#define trav(a, x) for (auto& a : x)
#define sz(x) (int)(x).size()
typedef pair<int, int> pii;
typedef vector<int> vi;
#pragma GCC diagnostic ignored "-Wsign-compare"
// util functions
int n, k;
VVI adj;
VII edg;
VI color;
VVI nodes;
VI dep;
VI par;
VI mx;
VI deg;
void init() {
adj.resize(n);
color.resize(n);
nodes.resize(k);
dep.resize(n);
mx.resize(n);
deg.assign(n, 0);
}
struct lca {
const int bits = 20;
VVI lft;
VI tin;
VI tout;
int tmr = 0;
void dfs(int u, int v) {
tin[u] = tmr++;
lft[u][0] = v;
for (int i = 1; i < bits; i++) {
int to = lft[u][i - 1];
if (to == -1)
lft[u][i] = -1;
else
lft[u][i] = lft[to][i - 1];
}
trav(x, adj[u]) {
if (x == v) continue;
dfs(x, u);
}
tout[u] = tmr++;
}
void init() {
tin.resize(n);
tout.resize(n);
lft.resize(n);
trav(x, lft) x.resize(bits);
dfs(0, 1);
}
bool isAnc(int u, int v) {
if (u == -1) return true;
if (v == -1) return false;
return tin[u] <= tin[v] and tout[u] >= tout[v];
}
int getLCA(int u, int v) {
if (isAnc(u, v)) return u;
if (isAnc(v, u)) return v;
for (int i = bits - 1; i >= 0; i--) {
if (!isAnc(lft[u][i], v)) {
u = lft[u][i];
}
}
return lft[u][0];
}
int getKthAnc(int u, int k) {
for (int i = bits - 1; i >= 0; i--) {
if (k & (1 << i)) {
u = lft[u][i];
}
}
return u;
}
} lcaa;
void dfs_dep(int u, int v) {
if (v == -1) dep[u] = 0;
trav(x, adj[u]) {
if (x == v) continue;
dep[x] = dep[u] + 1;
dfs_dep(x, u);
}
}
void dfs_mx(int u, int v) {
trav(x, adj[u]) {
if (x == v) continue;
dfs_mx(x, u);
mx[u] = min(mx[u], mx[x]);
}
}
int main() {
// #ifndef ONLINE_JUDGE
// freopen("mergers.in", "r", stdin);
// freopen("mergers.out", "w", stdout);
// #endif
cin.sync_with_stdio(0);
cin.tie(0);
cin >> n >> k;
if (n == 1) {
cout << "0\n";
return 0;
}
init();
for (int i = 0; i < n - 1; i++) {
int u, v;
cin >> u >> v;
u--, v--;
adj[u].PB(v);
adj[v].PB(u);
edg.PB(MP(u, v));
}
for (int i = 0; i < n; i++) {
int col;
cin >> col;
col--;
nodes[col].PB(i);
color[i] = col;
}
dfs_dep(0, -1);
mx = dep;
lcaa.init();
for (int i = 0; i < k; i++) {
int pt = nodes[i][0];
for (int j = 1; j < sz(nodes[i]); j++) {
int u = pt;
int v = nodes[i][j + 1];
int lca = lcaa.getLCA(u, v);
pt = lca;
int val = dep[lca];
}
for (int j = 0; j < sz(nodes[i]); j++) {
mx[nodes[i][j]] = min(mx[nodes[i][j]], dep[pt]);
}
}
dfs_mx(0, -1);
trav(x, edg) {
int u = lcaa.getKthAnc(x.x, dep[x.x] - mx[x.x]);
int v = lcaa.getKthAnc(x.y, dep[x.y] - mx[x.y]);
if (u != v) {
deg[u]++;
deg[v]++;
}
}
int ret = 0;
for (int i = 0; i < n; i++) {
if (deg[i] == 1) ret++;
}
cout << (ret + 1) / 2 << '\n';
return 0;
}
Compilation message
mergers.cpp: In function 'int main()':
mergers.cpp:196:11: warning: unused variable 'val' [-Wunused-variable]
196 | int val = dep[lca];
| ^~~
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
121 ms |
21552 KB |
Output is correct |
2 |
Correct |
123 ms |
26588 KB |
Output is correct |
3 |
Runtime error |
6 ms |
1900 KB |
Execution killed with signal 11 |
4 |
Halted |
0 ms |
0 KB |
- |
# |
Verdict |
Execution time |
Memory |
Grader output |
1 |
Correct |
1 ms |
364 KB |
Output is correct |
2 |
Incorrect |
1 ms |
364 KB |
Output isn't correct |
3 |
Halted |
0 ms |
0 KB |
- |