Submission #413836

#	Time	Username	Problem	Language	Result	Execution time	Memory
413836		rqi	Capital City (JOI20_capital_city)	C++14	100 / 100	848 ms	38084 KiB

capital_city

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

typedef long long ll;
typedef vector<int> vi;

#define pb push_back

#define sz(x) (int)(x).size()

void ckmin(int& a, int b){
	a = min(a, b);
}

const int MOD = 1e9+7;

#include <ext/pb_ds/assoc_container.hpp>
using namespace __gnu_pbds;

template<class T> using Tree = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

struct HashFunc{
	size_t operator()(ll a){
		return ((a+123485122314LL)^((a+14123901283002LL)>>30));
	}
};

gp_hash_table<ll, vi> g;

const int mx = 200005;
vi adj[mx];
int C[mx];
vi invC[mx];
int ans = MOD;

bool done[mx];

int sub[mx];

void genSub(int node, int p = -1){
	sub[node] = 1;
	for(auto u: adj[node]){
		if(done[u]) continue;
		if(u == p) continue;
		genSub(u, node);
		sub[node]+=sub[u];
	}
}

bool in_subtree[mx];
bool comp_inc[mx];

void setInSubtree(int node, bool val, int p = -1){
	in_subtree[node] = val;
	for(auto u: adj[node]){
		if(done[u]) continue;
		if(u == p) continue;
		setInSubtree(u, val, node);
	}
}

int par[mx];

void genPar(int node, int p = -1){
	par[node] = p;
	for(auto u: adj[node]){
		if(done[u]) continue;
		if(u == p) continue;
		genPar(u, node);
	}
}

void centroidDecomp(int node){
	// cout << "centroidDecomp: " << node << "\n";
	// cout.flush();
	genSub(node);

	auto genCen = [&](int node){
		int cen = node;
		int p = -1;
		while(true){
			bool found = 0;
			for(auto u: adj[cen]){
				if(done[u]) continue;
				if(u == p) continue;
				if(sub[u]*2 > sub[node]){
					found = 1;
					p = cen;
					cen = u;
					break;
				}
			}
			if(!found){
				break;
			}
		}
		return cen;
	};

	int cen = genCen(node);

	// cout << "foundCen: " << cen << "\n";
	// cout.flush();

	setInSubtree(cen, 1);
	genPar(cen);

	//solve
	vi comps;

	queue<int> q;
	comps.pb(C[cen]);
	comp_inc[C[cen]] = 1;
	q.push(C[cen]);

	

	bool bad = 0;
	while(sz(q)){
		int c = q.front();
		q.pop();
		for(auto u: invC[c]){
			if(!in_subtree[u]){
				bad = 1;
				break;
			}
		}
		if(bad) break;

		for(auto u: invC[c]){
			int p = par[u];
			if(p == -1) continue;
			if(comp_inc[C[p]]) continue;
			comps.pb(C[p]);
			comp_inc[C[p]] = 1;
			q.push(C[p]);
		}
	}

	// cout << "comps: ";
	// for(auto u: comps){
	// 	cout << u << " ";
	// }
	// cout << "\n";

	if(!bad){
		ckmin(ans, sz(comps)-1);
	}
	
	setInSubtree(cen, 0);
	for(auto u: comps){
		comp_inc[u] = 0;
	}

	done[cen] = 1;
	for(auto u: adj[cen]){
		if(done[u]) continue;
		centroidDecomp(u);
	}
}

int main(){
	cin.tie(0)->sync_with_stdio(0);
	int N, K;
	cin >> N >> K;
	for(int i = 1; i <= N-1; i++){
		int a, b;
		cin >> a >> b;
		adj[a].pb(b);
		adj[b].pb(a);
	}
	for(int i = 1; i <= N; i++){
		cin >> C[i];
		invC[C[i]].pb(i);
	}


	centroidDecomp(1);

	cout << ans << "\n";
}

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