Submission #1146936

#	Time	Username	Problem	Language	Result	Execution time	Memory
1146936		KK_1729	Capital City (JOI20_capital_city)	C++17	11 / 100	3097 ms	120704 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<vector<int>> graph;
vector<int> subtree;
vector<int> vis;
int ans = 200000;
vector<int> overall;
int counter = 0;
int dp(int x, int p = -1){
	if (vis[x]) return 0;
	// counter++;
	// cout << counter << endl;
	for (auto item: graph[x]){
		if (item == p) continue;
		assert(item != -1);
		subtree[x] += dp(item, x); 
	}
	return subtree[x];
}

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

vector<int> par;
vector<vector<int>> o;
vector<int> done;
vector<int> b;
vector<int> col;
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){
	// assert(!vis[x]);
	// cout << "m" << x << endl;
	dp(x);
	// cout << "d" << endl;
	int c = find_centroid(x, -1, subtree[x]);
	// cout <<"c" << 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] = 1;
	}
	// cout << "e" << endl;
	bool can = true;
	// if (e.size() == 1){
	// 	cout << c << endl;
	// 	cout << col[x] << col[c] << endl;
	// 	cout << o[col[x]].size() << overall[col[x]] << endl;
	// }
	for (auto item: e){
		if (o[item].size() != overall[item]) can = false;
	}
	// cout << "o" << c << e.size() << endl;
	if (can){
	// 	if (e.size() == 1){
	// 		cout << c << endl;
	// 		cout << col[c] << endl;
	// 	cout << overall[col[c]] << endl;
	// }
		ans = min(ans, ((int)e.size()-1));
	}
	// cout << "a" << endl;
	for (auto item: e){
		o[item].clear(); 
		done[item] = 0;
	}
	for (auto item: b) par[item] = -1;
	b.clear();
	// cout << "b" << endl;

	vis[c] = 1;
	for (auto node: graph[c]){
		if (!vis[node]) init_centroid(node, c);
	}
}
void solve(){
	int n, k; cin >> n >> k;

	graph.resize(n+100);
	par.resize(n+100);
	o.resize(n+100);
	overall.resize(n+100);
	done.resize(n+100);
	subtree.resize(n+100);
	vis.resize(n+100);
	col.resize(n+100);
	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