Submission #1320083

#	Time	Username	Problem	Language	Result	Execution time	Memory
1320083		JuanJL	Unique Cities (JOI19_ho_t5)	C++20	100 / 100	298 ms	61112 KiB

joi2019_ho_t5

#include <bits/stdc++.h>

#define fst first
#define snd second
#define pb push_back
#define forn(i,a,b) for(int i = a; i<b; i++)
#define SZ(x) (int)x.size()
#define ALL(x) x.begin(),x.end()
#define mset(a,v) memset(a,v,sizeof(a))
#define FIN ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);

using namespace std;
typedef long long ll;

const int MAXN = 200000+5;

ll n,m;
vector<ll> adj[MAXN];
ll typ[MAXN];
ll dist[MAXN];
ll len[MAXN];
ll res[MAXN];

void dfs(ll nd, ll p, ll dis){
	dist[nd]=dis;
	len[nd]=0;
	for(auto i:adj[nd]) if(i!=p){
		dfs(i,nd,dis+1);
		len[nd]=max(len[nd],len[i]+1);
	}
}

ll resact = 0;
ll cnt[MAXN];

stack<pair<ll,ll>> pila;

void ndfs(ll nd, ll p){
	vector<pair<ll,ll>> radj; 
	for(auto i:adj[nd]) if(i!=p) radj.pb({(len[i]+1)*-1,i});

	sort(ALL(radj));

	ll maxi = (SZ(radj)>1?radj[1].fst*-1:0);
	for(auto i:radj){
		while(!pila.empty() && pila.top().fst>=dist[nd]-maxi){
			cnt[typ[pila.top().snd]]--;
			if(cnt[typ[pila.top().snd]]==0) resact--;
			pila.pop();
		}

		pila.push({dist[nd],nd});
		cnt[typ[nd]]++;
		if(cnt[typ[nd]]==1) resact++;
		
		ndfs(i.snd,nd);

		if(!pila.empty() && pila.top().snd==nd){
			pila.pop();
			cnt[typ[nd]]--;
			if(cnt[typ[nd]]==0) resact--;
		}
		
		maxi=max(maxi,i.fst*-1);
	}
	

	while(!pila.empty() && pila.top().fst>=dist[nd]-maxi){
			cnt[typ[pila.top().snd]]--;
			if(cnt[typ[pila.top().snd]]==0) resact--;
			pila.pop();
		}

		//cout<<maxi<<'\n';
		//cout<<"pila ";
		//stack<pair<ll,ll>> rp = pila; while(!rp.empty()) cout<<"("<<rp.top().snd<<","<<rp.top().fst<<") ", rp.pop(); cout<<'\n';
	//cout<<"RES "<<nd<<" "<<dist[nd]<<" "<<resact<<'\n';
	res[nd]=max(res[nd],resact);
}

int main(){
	cin>>n>>m;

	ll u,v;
	forn(i,0,n-1){
		cin>>u>>v; u--; v--;
		adj[u].pb(v);
		adj[v].pb(u);
	}	

	forn(i,0,n) cin>>typ[i];

	dfs(0,-1,0);

	pair<ll,ll> d1 = {-1,-1}; forn(i,0,n) d1=max(d1,{dist[i],i});

	dfs(d1.snd,-1,0);
	pair<ll,ll> d2 = {-1,-1}; forn(i,0,n) d2=max(d2,{dist[i],i});

	mset(res,0);
	mset(cnt,0);

	//cout<<"desde "<<d1.snd<<'\n';

	ndfs(d1.snd,-1);

	dfs(d2.snd,-1,0);

	//cout<<"desde "<<d2.snd<<'\n';
	ndfs(d2.snd,-1);

	forn(i,0,n) cout<<res[i]<<'\n';


	return 0;
}

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