Submission #979221

#	Time	Username	Problem	Language	Result	Execution time	Memory
979221		happy_node	Capital City (JOI20_capital_city)	C++17	100 / 100	547 ms	41296 KiB

capital_city

#include <bits/stdc++.h>
using namespace std;
 
const int MX=2e5+5;
 
int N,K;
int A[MX],B[MX],S[MX];
 
vector<int> adj[MX],group[MX];
bool vis[MX], colVis[MX], subtree[MX], used[MX];
int sz[MX],par[MX];

int proc(int v) {
	queue<int> q;
	colVis[S[v]]=true;

	int res=1;

	for(auto u:group[S[v]]) {
		if(!subtree[u]) return K;
		vis[u]=true;
		q.push(u);
	}

	while(!q.empty()) {
		auto v=q.front(); q.pop();

		if(!colVis[S[v]]) {
			res+=1;
			for(auto u:group[S[v]]) {
				if(!subtree[u]) return K;
				vis[u]=true;
				q.push(u);
			}
			colVis[S[v]]=true;
		}

		if(par[v]!=0 && !vis[par[v]]) {
			if(!subtree[par[v]]) return K;
			vis[par[v]]=true;
			q.push(par[v]);
		}
	}

	return res;
}

int ans;

int getSize(int v, int p) {
	sz[v]=1;
	for(auto u:adj[v]) {
		if(u==p || used[u]) continue;
		sz[v]+=getSize(u,v); 
	}
	return sz[v];
}

int getCen(int v, int p, int s) {
	for(auto u:adj[v]) {
		if(u==p || used[u]) continue;
		if(2*sz[u]>s) return getCen(u,v,s);
	}
	return v;
}

vector<int> nodes;

void dfs(int v, int p) {
	par[v]=p;
	nodes.push_back(v);
	for(auto u:adj[v]) {
		if(u==p || used[u]) continue;
		dfs(u,v);
	}
}

void cdc(int v) {
	int s=getSize(v,v);
	int cen=getCen(v,v,s);
	used[cen]=true;

	nodes.push_back(cen);
	par[cen]=0;
	for(auto u:adj[cen]) {
		if(used[u]) continue;
		dfs(u,cen);
	}

	for(auto u:nodes) {
		subtree[u]=true;
		colVis[S[u]]=false;
		vis[u]=false;
	}
	ans=min(ans,proc(cen));

	for(auto u:nodes) {
		subtree[u]=false;
		colVis[S[u]]=false;
		vis[u]=false;
	}
	nodes.clear();		
	
	for(auto u:adj[cen]) {
		if(used[u]) continue;
		cdc(u);
	}
}

int main() {
	cin.tie(0); ios_base::sync_with_stdio(0);
 
	cin>>N>>K;
	for(int i=0;i<N-1;i++) {
		cin>>A[i]>>B[i];
		adj[A[i]].push_back(B[i]);
		adj[B[i]].push_back(A[i]);
	}
	for(int i=1;i<=N;i++) {
		cin>>S[i];
		group[S[i]].push_back(i);
	}
	ans=K;
	cdc(1);

	cout<<ans-1<<'\n';
}

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