Submission #827646

#	Time	Username	Problem	Language	Result	Execution time	Memory
827646		happypotato	Mergers (JOI19_mergers)	C++17	100 / 100	885 ms	136704 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
#define pii pair<int, int>
#define ff first
#define ss second
#define pb push_back
const int mxN = 5e5 + 1;
vector<int> adj[mxN];
int a[mxN];
vector<int> state[mxN];
int n, k;
void Input() {
	cin >> n >> k;
	for (int i = 0; i < n - 1; i++) {
		int u, v;
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for (int i = 1; i <= n; i++) {
		cin >> a[i];
		state[a[i]].pb(i);
	}
}
int tt = 0;
int dep[mxN];
pii range[mxN];
int lift[mxN][20];
void dfs1(int u = 1, int par = 0, int depth = 1) {
	dep[u] = depth;
	range[u].ff = ++tt;
	lift[u][0] = par;
	for (int &v : adj[u]) {
		if (v == par) continue;
		dfs1(v, u, depth + 1);
	}
	range[u].ss = tt;
}
void BuildLift() {
	for (int bit = 1; bit < 20; bit++) {
		for (int i = 1; i <= n; i++) {
			lift[i][bit] = lift[lift[i][bit - 1]][bit - 1];
		}
	}
}
int FindLCA(int u, int v) {
	if (dep[u] > dep[v]) swap(u, v);
	int diff = dep[v] - dep[u];
	for (int i = 19; i >= 0; i--) {
		if (diff & (1 << i)) {
			v = lift[v][i];
		}
	}
	if (u == v) return u;
	for (int i = 19; i >= 0; i--) {
		if (lift[u][i] != lift[v][i]) {
			u = lift[u][i];
			v = lift[v][i];
		}
	}
	return lift[u][0];
}
bool cmp(int &x, int &y) {
	return range[x].ff < range[y].ff;
}
int dp[mxN];
void dfs2(int u = 1, int par = 0) {
	for (int &v : adj[u]) {
		if (v == par) continue;
		dfs2(v, u);
		dp[u] += dp[v];
	}
}
int dsupp[mxN];
int DSUFind(int u) {
	if (dsupp[u] == 0) return u;
	return (dsupp[u] = DSUFind(dsupp[u]));
}
void DSUMerge(int u, int v) {
	u = DSUFind(u);
	v = DSUFind(v);
	if (u != v) dsupp[v] = u;
}
int deg[mxN];
void solve() {
	Input();
	dfs1();
	BuildLift();
	for (int i = 1; i <= k; i++) {
		sort(state[i].begin(), state[i].end(), cmp);
		for (int j = 1; j < (int)(state[i].size()); j++) {
			int u = state[i][j - 1], v = state[i][j];
			int lca = FindLCA(u, v);
			// colour edges from u to lca and from v to lca
			dp[u]++; dp[lca]--;
			dp[v]++; dp[lca]--;
		}
	}
	dfs2();
	for (int i = 2; i <= n; i++) {
		if (dp[i] > 0) {
			DSUMerge(i, lift[i][0]);
		}
	}
	// find degree of tree after merging
	for (int i = 2; i <= n; i++) {
		if (dp[i] == 0) {
			deg[DSUFind(i)]++;
			deg[DSUFind(lift[i][0])]++;
		}
	}
	int leaves = 0;
	for (int i = 1; i <= n; i++) {
		leaves += (deg[i] == 1);
	}
	cout << (leaves + 1) / 2 << endl;
}
int32_t main() {
	ios::sync_with_stdio(0); cin.tie(0);
	solve();
}

• Subtask 1: 10 / 10
• Subtask 2: 24 / 24
• Subtask 3: 14 / 14
• Subtask 4: 22 / 22
• Subtask 5: 30 / 30