Submission #518933

#	Time	Username	Problem	Language	Result	Execution time	Memory
518933		hivakarami	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	460 ms	49044 KiB

joi2019_ho_t5

#include<bits/stdc++.h>
 
using namespace std;

typedef long long int ll;
typedef long double ld;
 
#define f first
#define s second
#define pb push_back
#define pii pair<int, int>
#define mp make_pair
#define SZ(v) int(v.size())
#define all(v) (v).begin(),(v).end()
 
const int N = 2e5 + 100;
const ll inf = 1e16;
const ll base = 313;
const ll mod = 1009;
 
int h[N], mx[N], mx2[N];
vector<int> adj[N];
set<pii> st[N];
int ans[N], c[N], n, cnt = 0, Cnt[N];
deque<int> q;  
 
inline void add(int x)
{
	//cout << "add " << x << endl;
	cnt += (Cnt[c[x]] == 0); 
	Cnt[c[x]]++;
	q.push_back(x);
} 
 
 
inline void del()
{
	int u = q.back();
	//cout << "del " << u << endl;
	Cnt[c[u]]--;
	cnt -= (Cnt[c[u]] == 0);
	q.pop_back();
} 
 
void dfs(int u, int p)
{
	mx[u] = 0; mx2[u] = -1;
	for(auto x : adj[u])
	{
		if(x != p)
		{
			h[x] = h[u] + 1;
			dfs(x, u);
			if(mx[x] + 1 >= mx[u])
			{
				mx2[u] = mx[u];
				mx[u] = mx[x]+1;
			}
			else if(mx[x] + 1 >= mx2[u])
				mx2[u] = mx[x]+1;
		}
	}
}
 
inline pii find_root()
{
	dfs(0, -1);
	int r1 = 0, r2 = 0;
	for(int i = 0; i < n; i++)
	{
		if(h[i] > h[r1])
			r1 = i;
	}
	h[r1] = 0;
	dfs(r1, -1);
	for(int i = 0; i < n; i++)
	{
		if(h[i] > h[r2])
			r2 = i;
	//	cout << h[i] << ' ' << r2 << '\n';
	}
	
	return {r1, r2};
	
	
} 


void solve(int u, int p = -1)
{
		
	for(auto x : adj[u])
	{
		if(x != p && mx[x] + 1 == mx[u])
		{
			while(q.size() && h[u] - h[q.back()] <= mx2[u])
				del();
			
			add(u);
			solve(x, u);
		}
	}

	for(auto x : adj[u])
	{
		if(x != p && mx[x] + 1 != mx[u])
		{
			while(q.size() && h[u] - h[q.back()] <= mx[u])
				del();
			
			add(u);
			solve(x, u);
		}
	}
	
	while(q.size() && h[u] - h[q.back()] <= mx[u])
		del();
	
	/*		
	cout << "u " << u<< endl;
	for(int i = 0; i < 5; i++)
		cout << Cnt[i] << ' ';
	cout << endl;	
	cout << "/////////////////////////" << endl;		
	*/		
	ans[u] = max(ans[u], cnt);
		
}

int main()
{
    ios_base::sync_with_stdio(false), cin.tie(0), cout.tie(0);	


	int m;
	cin >> n >> m;
	for(int i = 0; i < n-1; i++)
	{
		int x, y;
		cin >> x >> y;
		x--; y--;
		adj[x].pb(y);
		adj[y].pb(x);
	}

	for(int i = 0; i < n; i++)
		cin >> c[i];

 	pii r = find_root();	
	
	
	//cout << r.f << ' ' << r.s << endl;
	
	h[r.f] = 0;
	dfs(r.f, -1);
	solve(r.f);
	
	h[r.s] = 0;
	dfs(r.s, -1);
	solve(r.s);
	
	
	for(int i = 0; i < n; i++)
		cout << ans[i] << '\n';
	
	
	
	
	return 0;
}
 
 
 
 
/*
5 4
1 2
2 3
3 4
3 5
1 2 1 2 4


*/
 
 

 
 

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