Submission #883042

#	Time	Username	Problem	Language	Result	Execution time	Memory
883042		8pete8	Capital City (JOI20_capital_city)	C++14	100 / 100	272 ms	100468 KiB

capital_city

#include<iostream>
#include<stack>
#include<map>
#include<vector>
#include<string>
#include<unordered_map>
#include <queue>
#include<cstring>
#include<limits.h>
#include<cmath>
#include<set>
#include<algorithm>
#include <iomanip>
#include<numeric>
#include<bitset>
using namespace std;
#define ll long long
#define f first
#define endl "\n"
#define s second
#define pii pair<int,int>
#define ppii pair<int,pii>
#define vi vector<int>
#define pb push_back
#define all(x) x.begin(),x.end()
#define rall(x) x.rbegin(),x.rend()
#define F(n) for(int i=0;i<n;i++)
#define lb lower_bound
#define ub upper_bound
#define fastio ios::sync_with_stdio(false);cin.tie(NULL);
#pragma GCC optimize ("03,unroll-loops")
using namespace std;
#define int long long
const int mod=1e9+7,mxn=2e5+5,lg=30,inf=1e18,minf=-1e18;
int n,k;
int h[mxn+10],up[mxn+10][lg+5],g[mxn+10],gl[mxn+10],low[mxn+10],disc[mxn+10];
vector<int>adj[mxn+10],group[mxn+10],adj2[mxn+10];
void dfs(int cur,int p){
	for(auto i:adj[cur]){
		if(i==p)continue;
		h[i]=h[cur]+1;
		up[i][0]=cur;
		dfs(i,cur);
	}
}
int lca(int a,int b){
	if(h[a]<h[b])swap(a,b);
	int k=h[a]-h[b];
	for(int i=0;i<=lg;i++)if(k&(1<<i))a=up[a][i];
	if(a==b)return a;
	for(int i=lg;i>=0;i--)if(up[a][i]!=up[b][i])a=up[a][i],b=up[b][i];
	return up[a][0];
}
bitset<mxn+10>on,vis;
stack<int>st;
int t=0,ans=inf;
int comp[mxn+10];
int compcnt=0;
void tarjan(int cur){
	vis[cur]=true;
	low[cur]=disc[cur]=++t;
	st.push(cur);
	on[cur]=true;
	for(auto i:adj2[cur]){
		if(!disc[i])tarjan(i),low[cur]=min(low[cur],low[i]);
		else if(on[i])low[cur]=min(low[cur],low[i]);
	}
	if(disc[cur]==low[cur]){
		compcnt++;
		bool yes=true;
		on[cur]=false;
		comp[cur]=compcnt;
		vector<int>v;
		int cnt=1;
		v.pb(cur);
		while(!st.empty()&&st.top()!=cur){
			v.pb(st.top());
			if(on[st.top()])cnt++;
			comp[st.top()]=compcnt;
			on[st.top()]=false;
			st.pop();
		}
		if(!st.empty())st.pop();
		for(auto i:v)for(auto j:adj2[i])yes&=(compcnt==comp[j]);
		if(yes)ans=min(ans,cnt-1);
	}
}
int32_t main(){
	fastio
	cin>>n>>k;
	for(int i=0;i<n-1;i++){
		int a,b;cin>>a>>b;
		adj[a].pb(b);
		adj[b].pb(a);
	}
	dfs(1,-1);
	for(int i=1;i<=n;i++)cin>>g[i],group[g[i]].pb(i);
	for(int i=1;i<=lg;i++)up[1][i]=1;
	for(int i=1;i<=lg;i++)for(int j=1;j<=n;j++)up[j][i]=up[up[j][i-1]][i-1];
	for(int i=1;i<=k;i++){
		if(group[i].empty())continue;
		int node=group[i][0];
		for(int j=1;j<group[i].size();j++)node=lca(node,group[i][j]);
		gl[i]=node;
	}
	for(int i=1;i<=n;i++)if(g[up[i][0]]!=g[i]&&i!=gl[g[i]])adj2[g[i]].pb(g[up[i][0]]);
	for(int i=1;i<=k;i++)if(group[i].size()&&!vis[i])tarjan(i);
	cout<<((ans==inf)?0:ans);
	return 0;
}

Compilation message (stderr)

capital_city.cpp: In function 'int32_t main()':
capital_city.cpp:103:16: warning: comparison of integer expressions of different signedness: 'long long int' and 'std::vector<long long int>::size_type' {aka 'long unsigned int'} [-Wsign-compare]
  103 |   for(int j=1;j<group[i].size();j++)node=lca(node,group[i][j]);
      |               ~^~~~~~~~~~~~~~~~

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