Submission #1143948

#TimeUsernameProblemLanguageResultExecution timeMemory
1143948monkey133Unique Cities (JOI19_ho_t5)C++20
100 / 100
536 ms56608 KiB
#include <bits/stdc++.h>
//#include "includeall.h"
//#include <ext/pb_ds/assoc_container.hpp>
//#include <ext/pb_ds/tree_policy.hpp>
#define endl '\n'
#define f first
#define s second
#define pb push_back
#define mp make_pair
#define lb lower_bound
#define ub upper_bound
#define input(x) scanf("%lld", &x);
#define input2(x, y) scanf("%lld%lld", &x, &y);
#define input3(x, y, z) scanf("%lld%lld%lld", &x, &y, &z);
#define input4(x, y, z, a) scanf("%lld%lld%lld%lld", &x, &y, &z, &a);
#define print(x, y) printf("%lld%c", x, y);
#define show(x) cerr << #x << " is " << x << endl;
#define show2(x,y) cerr << #x << " is " << x << " " << #y << " is " << y << endl;
#define show3(x,y,z) cerr << #x << " is " << x << " " << #y << " is " << y << " " << #z << " is " << z << endl;
#define all(x) x.begin(), x.end()
#define discretize(x) sort(x.begin(), x.end()); x.erase(unique(x.begin(), x.end()), x.end());
#define FOR(i, x, n) for (ll i =x; i<=n; ++i) 
#define RFOR(i, x, n) for (ll i =x; i>=n; --i) 
using namespace std;
mt19937_64 rnd(chrono::steady_clock::now().time_since_epoch().count());
//using namespace __gnu_pbds;
//#define ordered_set tree<int, null_type, less<int>, rb_tree_tag, tree_order_statistics_node_update>
//#define ordered_multiset tree<int, null_type, less_equal<int>, rb_tree_tag, tree_order_statistics_node_update>
typedef long long ll;
typedef long double ld;
typedef pair<ld, ll> pd;
typedef pair<string, ll> psl;
typedef pair<ll, ll> pi;
typedef pair<ll, pi> pii;
typedef pair<pi, pi> piii;
mt19937 rng(chrono::steady_clock::now().time_since_epoch().count());

ll n, m;
vector<ll> adj[200005];
ll dist[200005], maxdep[200005], ans[200005], color[200005];

// count unique element
ll freq[200005];
ll sum = 0;
deque<ll> dq;
void insert(ll x)
{
	dq.pb(x);
	x = color[x];
	freq[x]++;
	if (freq[x]==1) sum++;
}
void remove()
{
	ll x = color[dq.back()];
	dq.pop_back();
	freq[x]--;
	if (freq[x]==0) sum--;
}


// find diameter

pi dfs1(ll x = 1, ll p = -1, ll dis = 0)
{
	maxdep[x] = 0, dist[x] = 0;
	pi ret = mp(dis, x); 
	for (ll u: adj[x]) if (u!=p)  ret = max(ret, dfs1(u, x, dis + 1));
	return ret;
}

// precompute distance & maximum depth

void dfs2(ll x, ll p = -1)
{
	maxdep[x] = 1;
	for (ll u: adj[x]) if (u!=p) 
	{
		dist[u] = dist[x] + 1;
		dfs2(u, x);
		maxdep[x] = max(maxdep[x], maxdep[u] + 1);
	}
}

// calculate

void dfs3(ll x ,ll p = -1)
{
	// hld decomp
	vector<pi> nxt;
	for (ll u: adj[x]) if (u!=p) nxt.pb(mp(maxdep[u], u));
	sort(all(nxt), greater<pi> ());
	// go to deepest first, for monotonicity
	for (ll i=0; i<nxt.size(); ++i)
	{
		ll ban = nxt[0].f;
		if (i==0) ban = (nxt.size()==1)?0:nxt[1].f;
		while (dq.size() && dist[dq.back()] >= dist[x] - ban) remove();
		insert(x);
		dfs3(nxt[i].s, x);
		if (dq.size() && dq.back()==x) remove();
	}
	ll ban = (nxt.size())?nxt[0].f:0;
	while (dq.size() && dist[dq.back()] >= dist[x] - ban) remove();
	ans[x] = max(ans[x], sum);
	
	
}

int main()
{
	cin >> n >> m;
	for (ll i=0; i<n-1; ++i)
	{
		ll u, v;
		cin >> u >> v;
		adj[u].pb(v);
		adj[v].pb(u);
	}
	for (ll i=1; i<=n; ++i) cin >> color[i];
	ll start = dfs1().s;
	dfs2(start);
	dfs3(start);
	ll end = dfs1(start).s;
	dfs2(end);
	dfs3(end);
	for (ll i=1; i<=n; ++i) cout << ans[i] << endl;
	return 0;
}

#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...
#Verdict Execution timeMemoryGrader output
Fetching results...