Submission #360018

#	Time	Username	Problem	Language	Result	Execution time	Memory
360018		soroush	Mergers (JOI19_mergers)	C++14	100 / 100	2342 ms	141640 KiB

mergers

/*
#pragma GCC optimize("O2")
#pragma GCC optimize("Ofast")
#pragma GCC optimize("unroll-loops")
#pragma GCC target("avx,avx2,sse,sse2,fma")
*/
#include <bits/stdc++.h>

using namespace std;

typedef long long ll;
typedef long double ld;
typedef pair<int , int> pii;

mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

const int maxn = 5e5 + 100;
const ll mod = 1e9+7;
const ld PI = acos((ld)-1);

#define pb push_back
#define endl '\n'
#define dokme(x) cout << x , exit(0)
#define migmig ios::sync_with_stdio(false),cin.tie(0),cout.tie(0)
#define ms(x , y) memset(x , y , sizeof x)
ll pw(ll a, ll b, ll md = mod){ll res = 1;while(b){if(b&1){res=(a*res)%md;}a=(a*a)%md;b>>=1;}return(res);}

int n , k , c[maxn];
vector < int > adj[maxn] , col[maxn];
int root[maxn];

const int lg = 20;

struct LCA{
	int n;
	int par[lg+3][maxn] , h[maxn];
	void dfs(int v  , int p , vector < int > *adj){
		for(int u : adj[v]) if(u ^ p)
			h[u] = h[v] + 1 , par[0][u] = v , dfs(u , v , adj);
	}
	void init(int N , vector < int > *adj , int root = 1){
		n = N;
		dfs(root , 0 , adj);
		h[0] = -1;
		for(int j = 1 ; j <= lg ; j ++)
			for(int i = 1 ; i <= n ; i ++)
				par[j][i] = par[j - 1][par[j - 1][i]];
	}
	int lca(int u , int v){
		if(h[u] < h[v])
			swap(u , v);
		for(int i = lg ; i >= 0 ; i --)
			if(h[par[i][u]] >= h[v])
				u = par[i][u];
		if(u == v)
			return(u);
		for(int i = lg ; i >= 0 ; i --)
			if(par[i][u] ^ par[i][v])
				u = par[i][u] , v = par[i][v];
		return(par[0][v]);
	}
	int dist(int u , int v){
		return(h[u] + h[v] - 2*h[lca(u , v)]);
	}
};

LCA lca;

int par[maxn];

int getpar(int v){
	return((!par[v]) ? v : par[v] = getpar(par[v]));
}

void merge(int v , int u){
	if((u = getpar(u)) == (v = getpar(v)))
		return;
	par[u] = v;
}

int a[maxn];

void dfs(int v , int par = 0){
	for(int u : adj[v])if(u ^ par)
		dfs(u , v) , a[v] += a[u];
	if(a[v]) merge(v , par);
}

int d[maxn];

int32_t main(){
	migmig;
	cin >> n >> k;
	for(int i = 1 ; i < n ; i ++){
		int u , v;
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for(int i = 1 ; i <= n ; i ++)
		cin >> c[i] , col[c[i]].pb(i) , root[c[i]] = i;
	lca.init(n , adj);
	int ans = 0;
	for(int i = 1 ; i <= k ; i ++)
		for(int v : col[i])
			root[i] = lca.lca(root[i] , v);
	for(int i = 1 ; i <= k ; i ++)
		for(int v : col[i])
			a[v] ++ , a[root[i]] --;
	dfs(1);
	for(int i = 1 ; i <= n ; i ++){
		for(auto u : adj[i])
			if(getpar(i) ^ getpar(u))
				d[getpar(i)] ++ , d[getpar(u)] ++;
	}
	for(int i = 1 ; i <= n ; i ++)
		ans += (d[i] == 2);
	cout << (ans + 1) / 2;
	return(0);
}

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