Submission #218090

#	Time	Username	Problem	Language	Result	Execution time	Memory
218090		ryansee	Capital City (JOI20_capital_city)	C++14	100 / 100	1138 ms	41868 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());    //can be used by calling rng() or shuffle(A, A+n, rng)
inline long long rand(long long x, long long y) { return (rng() % (y+1-x)) + x; } //inclusivesss
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); }

typedef long long ll; 
typedef long double ld;
#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)
typedef pair<ll,ll>pi; typedef pair<ll,pi>spi; typedef pair<pi,pi>dpi;

#define LLINF ((long long)1e18)
#define INF int(1e9+1e6)
#define MAXN (200006)
ll n,k,C[MAXN],total,sz[MAXN],ans[MAXN],colour[MAXN],par[MAXN];
vector<int>v[MAXN],act[MAXN];
bitset<MAXN>done,vis;
void dfs1(int x,int p){
	++ total;
	sz[x]=1;
	for(auto i:v[x]) if(!done[i]) if(i^p) {
		dfs1(i,x);
		sz[x]+=sz[i];
	}
}
ll dfs2(int x,int p){
	for(auto i:v[x]) if(!done[i]&&i!=p&&sz[i]>total/2) return dfs2(i,x);
	return x;
}
void count(int x,int p){
	++ colour[C[x]];
	par[x]=p;
	for(auto i:v[x]) if(!done[i]) if(i^p) count(i,x);
}
void reset(int x,int p){
	colour[C[x]]=vis[C[x]]=0;
	for(auto i:v[x]) if(!done[i]) if(i^p) reset(i,x);
}
bool valid(ll c) {
	return colour[c] == siz(act[c]);
}
void dfs3(int x,int p){
	total=0;
	dfs1(x,p);
	x=dfs2(x,p);
	count(x,x);
	bool can=1;
	queue<ll> q;
	q.emplace(C[x]);
	vis[C[x]]=1;
	ans[x]=0;
	while(q.size()){
		ll c=q.front();
		q.pop();
		if(!valid(c)){
			can=0;
			break;
		}
		++ ans[x];
		for(auto i:act[c]) if(C[par[i]]^C[i]) if(!vis[C[par[i]]]) {
			q.emplace(C[par[i]]), vis[C[par[i]]]=1;
		}
	}
	if(!can) ans[x]=LLINF;
	else -- ans[x];
	reset(x,p);
	done[x]=1;
	for(auto i:v[x]) if(!done[i]) dfs3(i,x);
}
int main(){
	FAST
	cin>>n>>k;
	FOR(i,1,n-1){
		ll a,b;cin>>a>>b;
		v[a].eb(b), v[b].eb(a);
	}
	FOR(i,1,n)cin>>C[i],act[C[i]].pb(i);
	dfs3(1,-1);
	cout<<*min_element(ans+1,ans+n+1)<<'\n';
}

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