제출 #102402

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
102402		Dat160601	Mergers (JOI19_mergers)	C++17	100 / 100	2987 ms	204916 KiB

mergers

#include <bits/stdc++.h>
using namespace std;
#define mp make_pair
#define pb push_back
#define fi first
#define se second

const int N = 5e5 + 7;

int n, k, u, v, s[N], level[N], sub[N], par[20][N], cnt = 0, pset[N], ret[N], sz[N], deg[N], now[N], ans = 0;
bool vis[N], col[N];
vector <int> edge[N], tp[N], g[N], leaf;
unordered_map < long long, int > ed;

int read(){
	char p;
	while((p = getchar())){
		if(p < '0' || p > '9') continue;
		break;
	}
	int ret = p - '0';
	while((p = getchar())){
		if(p >= '0' && p <= '9'){
			ret *= 10;
			ret += p - '0';
		}
		else break;
	}
	return ret;
}

void predfs(int u, int p){
	int ch = 0;
	for(int v : edge[u]){
		if(v == p) continue;
		deg[u]++;
		par[0][v] = u;
		level[v] = level[u] + 1;
		predfs(v, u);
		ch++;
	}
	if(ch == 0) leaf.pb(u);
}

int lca(int u, int v){
	if(level[u] > level[v]) swap(u, v);
	for(int i = 19; i >= 0; i--) if(level[par[i][v]] >= level[u]) v = par[i][v];
	for(int i = 19; i >= 0; i--) if(par[i][u] != par[i][v]) u = par[i][u], v = par[i][v];
	if(u == v) return u;
	return par[0][u];
}

int fset(int x){
	if(pset[x] == x) return x;
	return pset[x] = fset(pset[x]);
}

void unionset(int u, int v){
	u = fset(u), v = fset(v);
	if(u == v) return;
	if(now[u] <= now[v]){
		pset[u] = v;
		sz[v] += sz[u];
		sz[u] = 0;
		now[v] += now[u];
		now[u] = 0;
	}
	else{
		pset[v] = u;
		sz[u] += sz[v];
		sz[v] = 0;
		now[u] += now[v];
		now[v] = 0;
	}
}

void dfs(int u, int p){
	int child = 0;
	for(int v : g[u]){
		if(v == p) continue;
		child ++;
		dfs(v, u);
	}
	if(child == 0 || (child == 1 && u == 1)) ans++;
}

int main(){
	n = read(), k = read();
	for(int i = 1; i < n; i++){
		u = read(), v = read();
		edge[u].pb(v);
		edge[v].pb(u);
	}
	for(int i = 1; i <= n; i++){
		s[i] = read();
		tp[s[i]].pb(i);
	}
	par[0][1] = 1;
	level[1] = 1;
	predfs(1, 1);
	for(int i = 1; i <= 19; i++){
		for(int j = 1; j <= n; j++){
			par[i][j] = par[i - 1][par[i - 1][j]];
		}
	}
	for(int i = 1; i <= k; i++){
		now[i] = 1;
		pset[i] = i;
		sz[i] = 1;
		if(tp[i].empty()) continue;
		int cur = tp[i][0];
		for(int j = 1; j < (int)tp[i].size(); j++){
			cur = lca(cur, tp[i][j]);
		}
		sub[cur]++;
	}
	queue <int> q;
	for(int x : leaf) q.push(x);
	while(!q.empty()){
		int u = q.front();
		q.pop();
		int p = fset(s[u]);
		sz[p] -= sub[u];
		deg[par[0][u]]--;
		if(!deg[par[0][u]]){
			q.push(par[0][u]);
		}
		if(sz[p] == 0) continue;
		unionset(p, s[par[0][u]]);
	}
	for(int i = 1; i <= k; i++){
		if(pset[i] == i) ret[i] = ++cnt;
	}
	for(int i = 1; i <= k; i++){
		if(pset[i] != i) ret[i] = ret[fset(i)];
	}
	for(int i = 1; i <= n; i++){
		int id = ret[s[i]];
		for(int j : edge[i]){
			int jd = ret[s[j]];
			if(id == jd) continue;
			if(ed[1LL * id * N + jd] > 0) continue;
			g[id].pb(jd);
			g[jd].pb(id);
			ed[1LL * id * N + jd]++;
			ed[1LL * jd * N + id]++;
		}
	}
	dfs(1, 1);
	if(ans == 1) printf("0");
	else printf("%d", (ans + 1) / 2);
}

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