제출 #1146939

#	제출 시각	아이디	문제	언어	결과	실행 시간	메모리
1146939		KK_1729	수도 (JOI20_capital_city)	C++20	11 / 100	3095 ms	61820 KiB

capital_city

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

// #define int long long
#define FOR(i,a,b) for (int i = (a); i < (b); ++i)
#define pb push_back
#define all(a) a.begin(), a.end()
// #define endl "\n"

void printVector(vector<int> a){
	for (auto x: a) cout << x << " ";
	cout << endl;
}
int min(int x, int y){
	if (x < y) return x;
	else return y;
}
vector<int> graph[200001];

int subtree[200001];
bool vis[200001];
int ans = 200000;
int overall[200001];
int counter = 0;
int dp(int x, int p = -1){
	if (vis[x]) return 0;
	for (int &item: graph[x]){
		if (item == p || vis[item]) continue;
		subtree[x] += dp(item, x); 
	}
	return subtree[x];
}

int find_centroid(int x, int p, int n){
	for (int &node: graph[x]){
		if (node == p) continue;
		if (!vis[node] && subtree[node]*2 > n) return find_centroid(node, x, n);
	}
	return x;
}

int par[200001];
vector<int> o[200001];
bool done[200001];
vector<int> b;

int col[200001];
void setup(int x, int p = -1){
	o[col[x]].pb(x);
	b.pb(x);
	for (auto node: graph[x]){
		if (node == p) continue;
		par[node] = x;
		setup(node, x);
	}
}
void init_centroid(int x = 1, int p = -1){
	// cout << x << endl;
	dp(x);
	int c = find_centroid(x, -1, subtree[x]);
	// cout << c << endl;
	par[c] = -1;
	setup(c);

	stack<int> s;
	s.push(col[c]);
	vector<int> e;
	while (!s.empty()){
		int current = s.top();
		s.pop();
		if (done[current]) continue;
		e.pb(current);
		for (auto item: o[current]){
			if (par[item] != -1 && par[item] != 0 && col[par[item]] != col[item] && !done[col[par[item]]]) s.push(col[par[item]]);
		}
		done[current] = true;
	}
	// cout << "e" << endl;
	bool can = true;

	for (auto item: e) if (o[item].size() != overall[item]) can = false;

	if (can) ans = min(ans, ((int)e.size()-1));
	

	for (auto item: e){
		o[item].clear(); 
		done[item] = false;
	}
	for (auto item: b) par[item] = -1;
	b.clear();

	vis[c] = true;
	for (auto node: graph[c]){
		if (!vis[node]) init_centroid(node, c);
	}
}
void solve(){
	int n, k; cin >> n >> k;
	ans = k-1;
	
	// graph.resize(n+1);
	// par.resize(n+1);
	// o.resize(n+1);
	// overall.resize(n+1);
	// done.resize(n+1);
	// subtree.resize(n+1);
	// vis.resize(n+1);
	// col.resize(n+1);
	
	
	FOR(i,0,n-1){
		int a, b; cin >> a >> b;
		graph[a].pb(b);
		graph[b].pb(a);
	}
	FOR(i,1,n+1){
		cin >> col[i];	
		overall[col[i]]++;
	}
	// cout << overall[75] << endl;
	// cout << endl;
	init_centroid();
	cout << ans << endl;

}


int32_t main(){
	ios::sync_with_stdio(false);cin.tie(nullptr);
	// freopen("in.in", "r", stdin);
	int t = 1; // cin >> t;
	while (t--) solve();
}

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