Submission #1048928

#	Time	Username	Problem	Language	Result	Execution time	Memory
1048928		1L1YA	Mergers (JOI19_mergers)	C++14	70 / 100	772 ms	262144 KiB

mergers

//1L1YA

#include<bits/stdc++.h>
using namespace std;

//#pragma GCC optimize ("O3,unrool-loops")
//#pragma GCC target("avx2,bmi,bmi2,lzcnt,popcnt")

typedef long long       ll;
typedef pair<ll,ll>     pll;
typedef pair<int,int>   pii;

#define Pb              push_back
#define F               first
#define S               second
#define endl            '\n'
#define sep             ' '
#define all(x)          x.begin(),x.end()
#define al(x,n)         x+1,x+n+1
#define SZ(x)           ((int)x.size())
#define lc              (id<<1)
#define rc              (lc|1)
#define mid             (l+r>>1)
#define dokme(x)        cout << x << endl, exit(0)
#define sik(x)          cout << x << endl;continue;
#define FastIO          ios::sync_with_stdio(0);cin.tie(0);cout.tie(0);
#define FileIO          freopen("input.txt","r",stdin);freopen("output.txt","w",stdout);

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

const ll oo=1e18;
const int mod=1e9+7;
const int inf=1e9+111;
const int N=5e5+11;
const int lg=20;

int n,k,ans,a[N],h[N],pr[N],cnt[N],par[lg][N];
vector<int> vc[N],adj[N];
set<int> st[N];
vector<pii> E;
bool mark[N];

void dfs(int v,int p=0){
	for(int u: adj[v])
		if(u!=p){
			h[u]=h[v]+1;
			par[0][u]=v;
			dfs(u,v);
		}
}

int jump(int v,int k){
	for(int i=0;i<lg;i++)
		if(k>>i & 1)
			v=par[i][v];
	return v;
}

int lca(int u,int v){
	if(h[u]>h[v])	swap(u,v);
	v=jump(v,h[v]-h[u]);
	if(u==v)	return v;
	for(int i=lg-1;~i;i--)
		if(par[i][u]!=par[i][v])
			u=par[i][u],v=par[i][v];
	return par[0][v];
}

void DFS(int v,int p=0){
	st[v].insert(a[v]);
	for(int u: adj[v])
		if(u!=p){
			DFS(u,v);
			cnt[v]+=cnt[u];
			if(SZ(st[v])<SZ(st[u]))	st[v].swap(st[u]);
			for(int x: st[u])	st[v].insert(x);
		}
	if(cnt[v]==SZ(st[v]))	mark[v]=1;
}

void dFs(int v,int p=0){
	for(int u: adj[v])
		if(u!=p){
			if(!mark[u])	pr[u]=pr[v];
			dFs(u,v);
		}
}

int main(){
	FastIO

	cin >> n >> k;
	for(int i=1;i<n;i++){
		int u,v;
		cin >> u >> v;
		adj[u].Pb(v);
		adj[v].Pb(u);
		E.Pb({u,v});
	}
	for(int i=1;i<=n;i++)	cin >> a[i],vc[a[i]].Pb(i);
	dfs(1);
	for(int i=1;i<lg;i++)
		for(int v=1;v<=n;v++)
			par[i][v]=par[i-1][par[i-1][v]];
	for(int i=1;i<=k;i++){
		int v=vc[i][0];
		for(int u: vc[i])	v=lca(u,v);
		cnt[v]++;
	}
	DFS(1);
	for(int i=1;i<=n;i++)	pr[i]=i;
	dFs(1);
	fill(cnt,cnt+n+1,0);
	for(auto [u,v]: E)
		if(pr[u]!=pr[v]){
			cnt[pr[u]]++;
			cnt[pr[v]]++;
			if(cnt[pr[u]]==1)	ans++;
			else if(cnt[pr[u]]==2)	ans--;
			if(cnt[pr[v]]==1)	ans++;
			else if(cnt[pr[v]]==2)	ans--;	
		}
	cout << (ans+1)/2 << endl;

	return 0;
}

Compilation message (stderr)

mergers.cpp: In function 'int main()':
mergers.cpp:114:11: warning: structured bindings only available with '-std=c++17' or '-std=gnu++17'
  114 |  for(auto [u,v]: E)
      |           ^

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