Submission #1314307

#	Time	Username	Problem	Language	Result	Execution time	Memory
1314307		vlomaczk	Capital City (JOI20_capital_city)	C++20	100 / 100	486 ms	49924 KiB

capital_city

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
typedef long long ll;
using namespace __gnu_pbds;
using namespace std;

template <typename T>
using ordered_set = tree<T, null_type, less<T>, rb_tree_tag, tree_order_statistics_node_update>;

ll inf = 1e18;
ll ans = inf;
ll M = 200'010;
vector<vector<ll>> g(M);
vector<ll> city(M), par(M), sajz(M), cnt(M), have(M), seen(M), is_off(M);
vector<vector<ll>> S(M);

void sajz_dfs(ll v, ll p) {
	sajz[v] = 1;
	par[v] = p;
	for(auto u : g[v]) {
		if(u==p || is_off[u]) continue;
		sajz_dfs(u,v);
		sajz[v] += sajz[u];
	}
}

ll find_centroid(ll v, ll ts) {
	for(auto u : g[v]) {
		if(is_off[u]) continue;
		if(u==par[v]) {
			if(ts-sajz[v] > ts/2) return find_centroid(u,ts);
		} else {
			if(sajz[u] > ts/2) return find_centroid(u,ts);
		}
	}
	return v;
}

void centroid_decomposition(ll v) {
	sajz_dfs(v,v);
	v = find_centroid(v,sajz[v]);
	sajz_dfs(v,v);
	queue<pair<ll, ll>> Q;	
	Q.push({v,0});
	vector<ll> used;
	while(Q.size()) {
		auto[x,p] = Q.front(); Q.pop();
		used.push_back(city[x]);
		cnt[city[x]]++;
		for(auto u : g[x]) {
			if(is_off[u] || u==p) continue;
			Q.push({u,x});
		}
	}
	ll res = 0;
	stack<ll> stos;
	stos.push(city[v]);
	while(stos.size()) {
		ll x = stos.top(); stos.pop();
		if(seen[x]) continue;
		seen[x] = 1;
		res++;
		if(cnt[x]!=have[x]) {
			res = inf;
			break;
		}
		for(auto y : S[x]) {
			stos.push(city[par[y]]);
		}
	}
	for(auto x : used) {
		seen[x] = 0;
		cnt[x] = 0;
	}
	ans = min(ans, res);
	is_off[v] = 1;
	for(auto u : g[v]) {
		if(is_off[u]) continue;
		centroid_decomposition(u);
	}
}

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

	ll n, r;
	cin >> n >> r;
	for(ll i=1; i<n; ++i) {
		ll a, b;
		cin >> a >> b;
		g[a].push_back(b);
		g[b].push_back(a);
	}
	for(ll i=1; i<=n; ++i) {
		cin >> city[i];
		have[city[i]]++;
		S[city[i]].push_back(i);
	}
	centroid_decomposition(1);
	cout << ans-1 << "\n";

	return 0;
}

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