Submission #216329

#	Time	Username	Problem	Language	Result	Execution time	Memory
216329		zscoder	Capital City (JOI20_capital_city)	C++17	11 / 100	3085 ms	57440 KiB

capital_city

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;
 
#define fi first
#define se second
#define mp make_pair
#define pb push_back
#define fbo find_by_order
#define ook order_of_key
 
typedef long long ll;
typedef pair<int,int> ii;
typedef vector<int> vi;
typedef long double ld; 
typedef tree<ll, null_type, less<ll>, rb_tree_tag, tree_order_statistics_node_update> pbds;

vi adj[222222];
int sub[222222];
int color[222222];
vector<set<int> > cities;
set<int> usedcolors;
int colorcnt[222222];
int ans=int(1e9);
int banned[222222]; //banned colors
bool cenvisited[222222];
set<int> cursub; //current subtree elements

void prep(int u, int p=-1)
{
	sub[u]=1;	
	colorcnt[color[u]]++;
	usedcolors.insert(color[u]);
	cursub.insert(u);
	for(int v:adj[u])
	{
		if(v==p) continue;
		if(cenvisited[v]) continue;
		prep(v,u);
		sub[u]+=sub[v];
	}
}

int centroid(int u, int p=-1, int r=-1)
{
	for(int v:adj[u])
	{
		if(v==p) continue;
		if(cenvisited[v]) continue;
		if(sub[v]*2>sub[r])
		{
			return centroid(v,u,r);
		}
	}
	return u;
}

bool valid(int c) //is a color valid in this subtree
{
	return (int(cities[c].size())==colorcnt[c]);
}

set<int> colorset;
set<int> processed;

void push(int c)
{
	if(colorset.find(c)==colorset.end()&&processed.find(c)==processed.end())
	{
		colorset.insert(c);
	}
}

int h[222222];
int par[222222];

void calch(int u, int p=-1)
{
	for(int v:adj[u])
	{
		if(v==p) continue;
		if(cenvisited[v]) continue;
		par[v]=u;
		h[v]=h[u]+1;
		calch(v,u);
	}
}


struct DSU
{
	int S;
	
	struct node
	{
		int p, rt;
	};
	vector<node> dsu;
	
	DSU(int n)
	{
		S = n;
		for(int i = 0; i < n; i++)
		{
			node tmp;
			tmp.p = i; tmp.rt = i;
			dsu.pb(tmp);
		}
	}
	void reset(int n)
	{
		dsu.clear();
		S = n;
		for(int i = 0; i < n; i++)
		{
			node tmp;
			tmp.p = i; tmp.rt = i;
			dsu.pb(tmp);
		}
	}
	
	void resetnode(int u)
	{
		dsu[u].p = u;
		dsu[u].rt = u;
	}
	
	int rt(int u)
	{
		if(dsu[u].p == u) return u;
		dsu[u].p = rt(dsu[u].p);
		return dsu[u].p;
	}
	
	void merge(int u, int v)
	{
		u = rt(u); v = rt(v);
		if(u == v) return ;
		if(rand()&1) swap(u, v);
		dsu[v].p = u;
		if(h[dsu[v].rt]<h[dsu[u].rt])
		{
			dsu[u].rt=dsu[v].rt;
		}
	}
	
	bool sameset(int u, int v)
	{
		if(rt(u) == rt(v)) return true;
		return false;
	}
	
	int getrt(int u)
	{
		return dsu[rt(u)].rt;
	}
};
DSU dsu(1);

void solve(int u)
{
	prep(u);
	int cent = centroid(u,-1,u);
	h[cent]=0; par[cent]=-1;
	calch(cent);
	//solve for this centroid
	colorset.insert(color[cent]);
	bool pos=1;
	while(!colorset.empty())
	{
		int c = (*colorset.begin()); processed.insert(c); colorset.erase(c);
		if(!valid(c)){pos=0; break;} //invalid color found
		//let's push all the stuff of this color
		for(int u:cities[c])
		{
			while(u!=cent)
			{
				u=dsu.getrt(u);
				if(u==cent) break;
				push(color[par[u]]);
				dsu.merge(par[u],u);
				u=par[u];
			}
		}
	}
	/*
	cerr<<"SOLVE WITH CENTROID = "<<cent<<'\n';
	cerr<<"POSSIBLE = "<<pos<<'\n';
	if(pos)
	{
		cerr<<"PROCESSED = ";
		for(int x:processed) cerr<<x<<' ';
		cerr<<'\n';
	}
	*/
	for(int x:cursub)
	{
		dsu.resetnode(x);
	}
	if(pos)
	{
		ans=min(ans,int(processed.size()));
	}
	cenvisited[cent]=1;
	for(int x:usedcolors) colorcnt[x]=0;
	usedcolors.clear();
	colorset.clear();
	processed.clear();
	cursub.clear();
	for(int v:adj[cent])
	{
		if(cenvisited[v]) continue;
		//recurse
		solve(v);
	}
}

int main()
{
	ios_base::sync_with_stdio(0); cin.tie(0);
	int n,k; cin>>n>>k;
	dsu.reset(n);
	for(int i=0;i<n-1;i++)
	{
		int u,v; cin>>u>>v; u--; v--;
		adj[u].pb(v); adj[v].pb(u);
	}
	cities.resize(k);
	for(int i=0;i<n;i++)
	{
		cin>>color[i]; color[i]--;
		cities[color[i]].insert(i);
	}
	solve(0);
	cout<<ans-1<<'\n'; //# of cities - 1
}

• Subtask 1: 1 / 1
• Subtask 2: 10 / 10
• Subtask 3: 0 / 30
• Subtask 4: 0 / 59