Submission #569280

#	Time	Username	Problem	Language	Result	Execution time	Memory
569280		hgj123	Capital City (JOI20_capital_city)	C++14	100 / 100	719 ms	42068 KiB

capital_city

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

#define FAST ios_base::sync_with_stdio(false); cin.tie(0);
#define pb push_back
#define eb emplace_back
#define ins insert
#define f first
#define s second
#define cbr cerr<<"hi\n"
#define mmst(x, v) memset((x), v, sizeof ((x)))
#define siz(x) ll(x.size())
#define all(x) (x).begin(), (x).end()
#define lbd(x,y) (lower_bound(all(x),y)-x.begin())
#define ubd(x,y) (upper_bound(all(x),y)-x.begin())
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());
inline long long rand(long long x, long long y) { return rng() % (y+1-x) + x; } //inclusive
string inline to_string(char c) {string s(1,c);return s;} template<typename T> inline T gcd(T a,T b){ return a==0?llabs(b):gcd(b%a,a); }

using ll=long long;
using ld=long double;
#define FOR(i,s,e) for(ll i=s;i<=ll(e);++i)
#define DEC(i,s,e) for(ll i=s;i>=ll(e);--i)
using pi=pair<ll,ll>; using spi=pair<ll,pi>; using dpi=pair<pi,pi>; 

long long LLINF = 1e18;
int INF = 1e9+1e6;
#define MAXN (200006)
int n, k, C[MAXN], ans, sz[MAXN], total, label[MAXN], vis[MAXN], co, p[MAXN];
vector<int> v[MAXN], c[MAXN];
bitset<MAXN> done;
void dfs1(int x,int p) {
	++ total, sz[x] = 1, label[x] = co;
	for(auto i:v[x]) if(i^p && !done[i]) dfs1(i, x), sz[x] += sz[i];
}
int dfs2(int x,int p) {
	for(auto i:v[x]) if(i^p && !done[i] && sz[i]>total/2) return dfs2(i, x);
	return x;
}
void hehe(int x,int P) {
	p[x] = P;
	for(auto i:v[x]) if(i^P && !done[i]) hehe(i, x);
}
void dfs3(int x) {
	total = 0, ++ co, dfs1(x, x);
	x = dfs2(x, x);
	hehe(x, x);
	int mi = 0;
	queue<int> q;
	for(auto i:c[C[x]]) {
		if(label[i] != co) {
			mi = INF;
			break;
		}
		vis[i]=co, q.emplace(i);
	}
	while(mi < ans && q.size()) {
		int x = q.front();
		q.pop();
		if(vis[p[x]] == co) continue;
		int y=p[x];
		++ mi;
		for(auto i:c[C[y]]) {
			if(label[i] != co) {
				mi = INF;
				break;
			}
			vis[i] = co, q.emplace(i);
		}
	}
	ans = min(ans, mi);
	done[x] = 1;
	for(auto i:v[x]) if(!done[i]) dfs3(i);
}
int main() {
	FAST
	cin>>n>>k;
	FOR(i,1,n-1) {
		int a,b;
		cin>>a>>b;
		v[a].eb(b), v[b].eb(a);
	}
	FOR(i,1,n) cin>>C[i], c[C[i]].eb(i);
	ans=k-1, dfs3(1);
	cout<<ans<<'\n';
}

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