Submission #519520

#	Time	Username	Problem	Language	Result	Execution time	Memory
519520		AriaH	Unique Cities (JOI19_ho_t5)	C++14	100 / 100	782 ms	75796 KiB

joi2019_ho_t5

/* I can do this all day */

#pragma GCC optimize("Ofast")

#include <bits/stdc++.h>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>

using namespace std;
using namespace __gnu_pbds;

typedef long long ll;
typedef long double ld;
typedef pair < int, int > pii;
typedef pair < ll, ll > pll;
typedef __int128 INT;

#define F first
#define S second
#define all(x) x.begin(),x.end()
#define Mp make_pair
#define point complex
#define endl '\n'
#define SZ(x) (int)x.size()
#define fast_io ios::sync_with_stdio(0), cin.tie(0), cout.tie(0)
#define file_io freopen("input.txt", "r+", stdin); freopen("output.txt", "w+", stdout);
#define mashtali return cout << "Hello, World!", 0;

template < typename T > using Tree = tree<T, null_type, less <T>, rb_tree_tag, tree_order_statistics_node_update>;

const int N = 2e5 + 10;
const int LOG = 20;
const ll mod = 1e9 + 7;
const ll inf = 8e18;
const double pi = acos(-1);
const ld eps = 1e-18;

ll pw(ll a, ll b, ll M, ll ret = 1) { if(a == 0) return 0; a %= M; while(b) { ret = (b & 1? ret * a % M : ret), a = a * a % M, b >>= 1; } return ret % M; }

mt19937_64 rng(chrono::steady_clock::now().time_since_epoch().count());

int n, m, fen[N], check[N], C[N], H[N], dp[N], mark[N], S[N], Ans[2][N], node[N];

void add(int i, int x)
{
	for(i += 3; i < N; i += i & -i) fen[i] += x;
}

int ask(int i, int ret = 0)
{
	if(i < 1) return 1;
	for(i += 3; i; i -= i & -i) ret += fen[i];
	return ret;
}

void upd(int l, int r, int x)
{
	l = max(1, l);
	if(l > r) return;
	add(l, x);
	add(r + 1, -x);
}

set < int > st;

vector < int > G[N];

void pre(int v, int P = 0)
{
	H[v] = H[P] + 1;
	dp[v] = 1;
	for(auto u : G[v])
	{
		if(u == P) continue;
		pre(u, v);
		dp[v] = max(dp[v], dp[u] + 1);
	}
}

bool cmp(int i, int j)
{
	return dp[i] < dp[j];
}

inline int Add(int v)
{
	if(!mark[C[v]])
	{
		mark[C[v]] = 1;
		st.insert({-H[v]});
		return 1;
	}
	return 0;
}

inline void Del(int v)
{
	if(st.count(-H[v]))
	{
		st.erase(-H[v]);
		mark[C[v]] = 0;
	}
}

void dfs(int v, int P, int tp)
{
	node[H[v]] = v;
	sort(all(G[v]), cmp);
	int Mx1 = 0, Mx2 = 0, id1 = -1;
	for(auto u : G[v])
	{
		if(u == P) continue;
		if(dp[u] > Mx1)
		{
			Mx2 = Mx1;
			Mx1 = dp[u];
			id1 = u;
		}
		else if(dp[u] > Mx2)
		{
			Mx2 = dp[u];
		}
		upd(H[v] - dp[u], H[v] - 1, 1);
	}
	for(auto u : G[v])
	{
		if(u == P || u == id1) continue;
		Ans[tp][u] = Ans[tp][v];
		upd(H[v] - dp[u], H[v] - 1, -1);
		dfs(u, v, tp);
		upd(H[v] - dp[u], H[v] - 1, 1);
	}
	if(id1 == -1)
	{
		Ans[tp][v] += Add(P);
		Del(P);
	}
	else
	{
		Ans[tp][id1] = Ans[tp][v];
		upd(H[v] - Mx1, H[v] - 1, -1);
		vector < int > vec;
		///printf("here v = %d id1 = %d Mx1 = %d get = %d %d\n", v, id1, Mx1, ask(H[v] - Mx1), ask(H[v] - Mx1 + 1));
		if(ask(H[v] - Mx1) == 0)
		{
			Ans[tp][id1] += Add(node[H[v] - Mx1]);
			vec.push_back(node[H[v] - Mx1]);
		}
		if(ask(H[v] - Mx1 + 1) == 0 && Mx1 > 1)
		{
			Ans[tp][id1] += Add(node[H[v] - Mx1 + 1]);
			vec.push_back(node[H[v] - Mx1 + 1]);
		}
		dfs(id1, v, tp);
		while(SZ(vec))
		{
			Del(vec.back());
			vec.pop_back();
		}
		upd(H[v] - Mx1, H[v] - 1, 1);
	}
	for(auto u : G[v])
	{
		if(u == P) continue;
		upd(H[v] - dp[u], H[v] - 1, -1);
	}
	///printf("v = %d Ans = %d\n", v, Ans[tp][v]);
}

int fir[N], sec[N];

int main()
{
	fast_io;
	cin >> n >> m;
	for(int i = 1; i < n; i ++)
	{
		int a, b;
		cin >> a >> b;
		G[a].push_back(b);
		G[b].push_back(a);
	}
	for(int i = 1; i <= n; i ++)
	{
		cin >> C[i];
	}
	pre(1);
	int L = 0, R = 0;
	for(int i = 1; i <= n; i ++)
	{
		if(H[i] > H[L]) L = i;
	}
	pre(L);
	for(int i = 1; i <= n; i ++)
	{
		fir[i] = H[i];
		if(H[i] > H[R]) R = i;
	}
	pre(R);
	for(int i = 1; i <= n; i ++)
	{
		sec[i] = H[i];
		if(fir[i] > sec[i])
		{
			S[i] = 0;
		}
		else
		{
			S[i] = 1;
		}
	}
	pre(L);
	dfs(L, 0, 0);
	st.clear();
	memset(mark, 0, sizeof mark);
	memset(check, 0, sizeof check);
	pre(R);
	dfs(R, 0, 1);
	for(int i = 1; i <= n; i ++)
	{
		cout << Ans[S[i]][i] << "\n";
	}

	return 0;
}

/* check corner case(n = 1?), watch for negetive index or overflow */

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