Submission #203565

#	Time	Username	Problem	Language	Result	Execution time	Memory
203565		kig9981	Unique Cities (JOI19_ho_t5)	C++17	100 / 100	1861 ms	110348 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
 
#ifdef NON_SUBMIT
#define TEST(n) (n)
#define tout cerr
#else
#define TEST(n) ((void)0)
#define tout cin
#endif
 
using namespace std;

vector<int> adj[200000], Q[400000], tree[400000];
int res, N, hld[200000], W[200000], depth[200000], parent[200000], P[400000][19], dist[400000], sdist[400000], ans[200000], C[200000], CC[400000], cnt[200001];

void dfs(int c)
{
	W[c]=1;
	for(auto n: adj[c]) if(W[n]==0) {
		parent[n]=c;
		depth[n]=depth[c]+1;
		dfs(n);
		W[c]+=W[n];
	}
}

void dfs2(int c)
{
	for(auto n: adj[c]) if(W[n]<W[c]) {
		hld[n]=2*W[n]>=W[c] ? hld[c]:n;
		dfs2(n);
	}
}

int LCA(int a, int b)
{
	while(hld[a]!=hld[b]) {
		if(depth[hld[a]]<depth[hld[b]]) swap(a,b);
		a=parent[hld[a]];
	}
	return depth[a]<depth[b] ? a:b;
}

int get_dist(int a, int b)
{
	b=b<N ? b:parent[b-N];
	return depth[a]+depth[b]-2*depth[LCA(a,b)];
}

void solve(int c)
{
	for(auto n: adj[c]) if(W[n]<W[c]) {
		solve(n);
		if(dist[c]<dist[n]+1) {
			sdist[c]=dist[c];
			dist[c]=dist[n]+1;
		}
		else if(sdist[c]<dist[n]+1) sdist[c]=dist[n]+1;
	}
	P[c][0]=c;
	for(auto n: adj[c]) if(W[n]<W[c] && sdist[c]<dist[n]+1)  {
		int np=n;
		for(int i=18;i>=0;i--) if(get_dist(c,P[np][i])<=sdist[c]) np=P[np][i];
		if(get_dist(c,np)>sdist[c]) P[c][0]=np;
		else if(get_dist(c,P[np][0])>sdist[c]) P[c][0]=P[np][0];
	}
	for(int j=1;j<19;j++) P[c][j]=P[P[c][j-1]][j-1];
}

void solve2(int c)
{
	int m=N+c, sm=N, tm=N;
	for(int j=1;j<19;j++) P[N+c][j]=P[P[N+c][j-1]][j-1];
	for(auto n: adj[c]) if(W[n]<W[c]) {
		if(dist[m]<dist[n]) {
			tm=sm;
			sm=m;
			m=n;
		}
		else if(dist[sm]<dist[n]) {
			tm=sm;
			sm=n;
		}
		else if(dist[tm]<dist[n]) tm=n;
	}
	for(auto n: adj[c]) if(W[n]<W[c]) {
		int np;
		dist[N+n]=dist[m]==dist[n] ? dist[sm]+1:dist[m]+1;
		if(m==n) {
			np=sm;
			for(int i=18;i>=0;i--) if(get_dist(c,P[np][i])<=dist[tm]+1) np=P[np][i];
			if(get_dist(c,np)>dist[tm]+1) P[N+n][0]=np;
			else if(get_dist(c,P[np][0])>dist[tm]+1) P[N+n][0]=P[np][0];
		}
		else if(sm==n) {
			np=m;
			for(int i=18;i>=0;i--) if(get_dist(c,P[np][i])<=dist[tm]+1) np=P[np][i];
			if(get_dist(c,np)>dist[tm]+1) P[N+n][0]=np;
			else if(get_dist(c,P[np][0])>dist[tm]+1) P[N+n][0]=P[np][0];
		}
		else {
			np=m;
			for(int i=18;i>=0;i--) if(get_dist(c,P[np][i])<=dist[sm]+1) np=P[np][i];
			if(get_dist(c,np)>dist[sm]+1) P[N+n][0]=np;
			else if(get_dist(c,P[np][0])>dist[sm]+1) P[N+n][0]=P[np][0];
		}
		solve2(n);
	}
}

void dfs3(int c)
{
	CC[c]=2;
	if(++cnt[C[c<N ? c:parent[c-N]]]==1) res++;
	for(auto i: Q[c]) ans[i]=res;
	for(auto n: tree[c]) if(CC[n]==0) dfs3(n);
	if(--cnt[C[c<N ? c:parent[c-N]]]==0) res--;
}

int main()
{
	ios::sync_with_stdio(false);
	cin.tie(NULL); cout.tie(NULL);
	TEST(freopen("input.txt","r",stdin));
	TEST(freopen("output.txt","w",stdout));
	TEST(freopen("debug.txt","w",stderr));
	int M;
	cin>>N>>M;
	memset(P,-1,sizeof(P));
	dist[P[N][0]=N]=-1;
	for(int i=1;i<N;i++) {
		int a, b;
		cin>>a>>b;
		adj[--a].push_back(--b);
		adj[b].push_back(a);
		P[N+i][0]=N+i;
	}
	dfs(0); dfs2(0);
	solve(0); solve2(0);
	for(int i=0;i<N;i++) {
		cin>>C[i];
		if(dist[i]>dist[N+i]+1) {
			int c=i;
			for(int j=18;j>=0;j--) if(get_dist(i,P[c][j])<=dist[N+i]+1) c=P[c][j];
			if(get_dist(i,c)>dist[N+i]+1) Q[c].push_back(i);
			else if(get_dist(i,P[c][0])>dist[N+i]+1) Q[P[c][0]].push_back(i);
		}
		else if(dist[i]<dist[N+i]+1) {
			int c=N+i;
			for(int j=18;j>=0;j--) if(get_dist(i,P[c][j])<=dist[i]) c=P[c][j];
			if(get_dist(i,c)>dist[i]) Q[c].push_back(i);
			else if(get_dist(i,P[c][0])>dist[i]) Q[P[c][0]].push_back(i);
		}
	}
	for(int i=0;i<2*N;i++) tree[P[i][0]].push_back(i);
	CC[N]=2;
	for(int i=0;i<2*N;i++) if(CC[i]==0) {
		int r=P[i][18];
		for(;CC[r]==0;r=P[r][0]) {
			CC[r]=1;
			if(++cnt[C[r<N ? r:parent[r-N]]]==1) res++;
		}
		for(;CC[r]==1;r=P[r][0]) dfs3(r);
		for(;CC[r]==2;r=P[r][0]) {
			CC[r]=3;
			if(--cnt[C[r<N ? r:parent[r-N]]]==0) res--;
		}
	}
	for(int i=0;i<N;i++) cout<<ans[i]<<'\n';
	return 0;
}

• Subtask 1: 4 / 4
• Subtask 2: 32 / 32
• Subtask 3: 32 / 32
• Subtask 4: 32 / 32