Submission #1048932

#	Time	Username	Problem	Language	Result	Execution time	Memory
1048932		1L1YA	Mergers (JOI19_mergers)	C++17	100 / 100	436 ms	122548 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];
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){
	for(int u: adj[v])
		if(u!=p){
			DFS(u,v);
			cnt[v]+=cnt[u];
		}
	if(!cnt[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[u]++;
		cnt[v]-=SZ(vc[i]);
	}
	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:6:40: warning: bad option '-funrool-loops' to pragma 'optimize' [-Wpragmas]
    6 | #pragma GCC optimize ("O3,unrool-loops")
      |                                        ^
mergers.cpp:42:23: warning: bad option '-funrool-loops' to attribute 'optimize' [-Wattributes]
   42 | void dfs(int v,int p=0){
      |                       ^
mergers.cpp:51:21: warning: bad option '-funrool-loops' to attribute 'optimize' [-Wattributes]
   51 | int jump(int v,int k){
      |                     ^
mergers.cpp:58:20: warning: bad option '-funrool-loops' to attribute 'optimize' [-Wattributes]
   58 | int lca(int u,int v){
      |                    ^
mergers.cpp:68:23: warning: bad option '-funrool-loops' to attribute 'optimize' [-Wattributes]
   68 | void DFS(int v,int p=0){
      |                       ^
mergers.cpp:77:23: warning: bad option '-funrool-loops' to attribute 'optimize' [-Wattributes]
   77 | void dFs(int v,int p=0){
      |                       ^
mergers.cpp:85:10: warning: bad option '-funrool-loops' to attribute 'optimize' [-Wattributes]
   85 | int main(){
      |          ^

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