Submission #542619

#	Time	Username	Problem	Language	Result	Execution time	Memory
542619		alvingogo	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	449 ms	39600 KiB

joi2019_ho_t5

#include <bits/stdc++.h>
#pragma GCC optimize("Ofast")
#define AquA cin.tie(0);ios_base::sync_with_stdio(0);
#define fs first
#define sc second
#define cd complex<double>
#define p_q priority_queue
using namespace std;

struct no{
	vector<int> ch;
	int co;
	int dis;
	int hh=0;
	int ans=0;
};
vector<no> v;
vector<int> cnt;
int cz=0;
int t=0;
void dfs1(int r,int f,int dis){
	if(t==-1 || dis>v[t].dis){
		t=r;
	}
	v[r].dis=dis;
	for(auto y:v[r].ch){
		if(y!=f){
			dfs1(y,r,dis+1);
		}
	}
}
stack<int> s;
void dfs2(int r,int f,int dis){
	v[r].hh=v[r].dis=dis;
	for(auto y:v[r].ch){
		if(y!=f){
			dfs2(y,r,dis+1);
			v[r].hh=max(v[r].hh,v[y].hh);
		}
	}
}
void add(int r){
	s.push(r);
	cnt[v[r].co]++;
	if(cnt[v[r].co]==1){
		cz++;
	}
}
void del(){
	int r=s.top();
	s.pop();
	cnt[v[r].co]--;
	if(cnt[v[r].co]==0){
		cz--;
	}
}
void dfs3(int r,int f){
	while(s.size() && v[s.top()].dis>v[r].dis){
		del();
	}
	int mx=0,mn=0;
	for(auto y:v[r].ch){
		if(y!=f){
			if(v[y].hh>mx){
				mn=mx;
				mx=v[y].hh;
			}
			else if(v[y].hh>mn){
				mn=v[y].hh;
			}
		}
	}
	int z=-1;
	for(auto y:v[r].ch){
		if(y!=f){
			if(v[y].hh==mx){
				z=y;
				while(s.size() && v[s.top()].dis>=v[r].dis-(mn-v[r].dis)){
					del();
				}
				add(r);
				dfs3(y,r);
				break;
			}
		}
	}
	while(s.size() && v[s.top()].dis>=v[r].dis-(mx-v[r].dis)){
		del();
	}
	for(auto y:v[r].ch){
		if(y!=f && y!=z){
			if(!s.size() || s.top()!=r){
				add(r);
			}
			dfs3(y,r);
		}
	}
	if(s.size() && s.top()==r){
		del();
	}
	v[r].ans=max(v[r].ans,cz);
}
int main(){
	AquA;
	int n,m;
	cin >> n >> m;
	v.resize(n);
	cnt.resize(m+1);
	for(int i=1;i<n;i++){
		int a,b;
		cin >> a >> b;
		a--;
		b--;
		v[a].ch.push_back(b);
		v[b].ch.push_back(a);
	}
	for(int i=0;i<n;i++){
		cin >> v[i].co;
	}
	int y=0;
	dfs1(y,y,0);
	int a=t;
	t=-1;
	dfs1(a,a,0);
	int b=t;
	for(int i=0;i<2;i++){
		dfs2(a,a,0);
		dfs3(a,a);
		a=b;
	}
	for(int i=0;i<n;i++){
		cout << v[i].ans << "\n";
	}
	return 0;
}

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