/*
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, bool deb = false) {
for (int i = bits - 1; i >= 0; i--) {
if (k & (1 << i)) {
if (deb) debug(u, lft[u][i], 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]);
}
}
void checking(int u, int v, VI& vals) {
for (auto& x : adj[u]) {
if (x == v) continue;
if (vals[x] != vals[u] and lcaa.isAnc(vals[x], vals[u])) assert(false);
}
trav(x, adj[u]) {
if (x != v) checking(x, u, vals);
}
}
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];
int lca = lcaa.getLCA(u, v);
pt = 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++;
}
VI vals(n);
for (int i = 0; i < n; i++) {
vals[i] = lcaa.getKthAnc(i, dep[i] - mx[i]);
}
cout << (ret + 1) / 2 << '\n';
checking(0, -1, vals);
return 0;
}
Compilation message
mergers.cpp: In member function 'int lca::getKthAnc(int, int, bool)':
mergers.cpp:19:20: warning: statement has no effect [-Wunused-value]
19 | #define debug(...) 42
| ^~
mergers.cpp:127:18: note: in expansion of macro 'debug'
127 | if (deb) debug(u, lft[u][i], i);
| ^~~~~
| # |
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 |
123 ms |
22000 KB |
Output is correct |
| 2 |
Correct |
123 ms |
26992 KB |
Output is correct |
| 3 |
Correct |
3 ms |
1004 KB |
Output is correct |
| 4 |
Correct |
3 ms |
1004 KB |
Output is correct |
| 5 |
Correct |
1 ms |
364 KB |
Output is correct |
| 6 |
Correct |
1 ms |
364 KB |
Output is correct |
| 7 |
Correct |
3 ms |
1004 KB |
Output is correct |
| 8 |
Incorrect |
261 ms |
25824 KB |
Output isn't correct |
| 9 |
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 |
- |
