Submission #220391

#	Time	Username	Problem	Language	Result	Execution time	Memory
220391		Andrei_Cotor	Capital City (JOI20_capital_city)	C++11	100 / 100	812 ms	35560 KiB

capital_city

#include<iostream>
#include<vector>
#include<deque>
#include<algorithm>

using namespace std;

vector<int> Adj[200005];
vector<int> Towns[200005];
vector<int> RelevantCities;
bool Centroid[200005];
bool Merged[200005];
int Sz[200005],P[200005];
int NrCity[200005],City[200005];
int rez;

void getsz(int nod, int p)
{
	Sz[nod]=1;
	Towns[City[nod]].push_back(nod);
	RelevantCities.push_back(City[nod]);

	for(auto other:Adj[nod])
	{
		if(other==p || Centroid[other])
			continue;

		getsz(other,nod);

		Sz[nod]+=Sz[other];
	}
}

void getParents(int nod, int p)
{
	P[nod]=p;
	for(auto other:Adj[nod])
	{
		if(other==p || Centroid[other])
			continue;

		getParents(other,nod);
	}
}

int getCentroid(int nod, int p, int sztot)
{
	int nxt=-1;
	for(auto other:Adj[nod])
	{
		if(other==p || Centroid[other])
			continue;

		if(sztot/2<Sz[other])
			nxt=other;
	}

	if(nxt!=-1)
		return getCentroid(nxt,nod,sztot);
	return nod;
}

void centroidDecomp(int nod)
{
	getsz(nod,0);

	nod=getCentroid(nod,0,Sz[nod]);
	Centroid[nod]=true;

	getParents(nod,0);

	deque<int> Q;
	for(auto town:Towns[City[nod]])
	{
		if(town==nod)
			continue;
		Q.push_back(town);
	}

	Merged[City[nod]]=1;
	while(!Q.empty())
	{
		int town=Q.front();
		Q.pop_front();

		if(!Merged[City[P[town]]])
		{
			Merged[City[P[town]]]=1;
			for(auto other:Towns[City[P[town]]])
				Q.push_back(other);
		}
	}

	int nrmerges=0;
	bool ok=true;
	for(auto city:RelevantCities)
	{
		if(Merged[city])
		{
			if((int)Towns[city].size()!=NrCity[city])
				ok=false;

			nrmerges++;
			Merged[city]=0;
		}
		Towns[city].clear();
	}
	RelevantCities.clear();

	nrmerges--;
	if(ok)
		rez=min(rez,nrmerges);

	for(auto other:Adj[nod])
	{
		if(!Centroid[other])
			centroidDecomp(other);
	}
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(0);
	cout.tie(0);

	int n,k;
	cin>>n>>k;

	for(int i=1; i<n; i++)
	{
		int x,y;
		cin>>x>>y;
		Adj[x].push_back(y);
		Adj[y].push_back(x);
	}

	for(int i=1; i<=n; i++)
	{
		cin>>City[i];
		NrCity[City[i]]++;
	}

	rez=1000000000;
	centroidDecomp(1);

	cout<<rez<<"\n";
	return 0;
}

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