LOJ#3014. 「JOI 2019 Final」独特的城市

首先找出树的直径,令它的端点分别为\(S,T\).

先对\(S\)进行处理.

第一步,以\(S\)为根节点进行长链剖分.

对于每个点\(u\)，当走重儿子时，求出所有轻儿子的子树中的最长链长度,即次长链\(L\).

显然对于点\(u\)的重儿子,它上面\(L\)个节点就不是独特的城市了.

此时我们可以开一个全局的栈记录该点到根路径上可能作为答案的点.

此时就将它上面不合法的弹出.

然后是特产的处理.

我们可以开一个全局桶,\(cnt_i\)储存下当前栈内特产\(i\)的个数.

弹出时就将\(cnt[s[top]]--\),如果变为\(0\)了,就\(res--\),加入时类似.

对于\(u\)的轻儿子,它上面\(dep[son[x]]\),即重儿子的\(L\),就不是独特的城市了,也要将其弹出.

注意要在弹栈之后再把\(u\)压入栈中，而且遍历每个儿子前都要重新将\(u\)压入栈中(想一想,为什么).

又因为栈中只记录该点到根路径上可能作为答案的点,所以我们再以\(T\)为根做一遍,答案取\(max\)即可.

#pragma GCC optimize(3)
#include<bits/stdc++.h>
#define il inline
#define rg register
#define gi read<int>
#define File(x) freopen(x".in","r",stdin);freopen(x".out","w",stdout)
using namespace std;
const int O = 2e5 + 10;
struct Edge{int to, nt;}e[O << 1];
template<class TT>
il TT read() {
	TT o = 0,fl = 1; char ch = getchar();
	while (!isdigit(ch) && ch != '-') ch = getchar();
	if (ch == '-') fl = -1, ch = getchar();
	while (isdigit(ch)) o = o * 10 + ch - '0', ch = getchar();
	return fl * o;
}
int n, m, head[O], c[O], cnt;
il void add(int u, int v) {
	e[++cnt] = (Edge){v, head[u]};
	head[u] = cnt;
}
namespace cpp1{
	int dep[O], cnt[O], ans;
	bool vis[O];
	il void dfs(int u, int fa) {
		dep[u] = dep[fa] + 1;
		for (int i = head[u]; i; i = e[i].nt) {
			int v = e[i].to;
			if (v == fa) continue;
			dfs(v, u);
		}
	}
	il void main() {
		for (int i = 1; i <= n; ++i) {
			dfs(i, 0); ans = 0;
			memset(cnt, 0, sizeof cnt); memset(vis, 0, sizeof vis);
			cnt[1] = 1;
			for (int j = 1; j <= n; ++j) ++cnt[dep[j]];
			for (int j = 1; j <= n; ++j) if (cnt[dep[j]] == 1) vis[c[j]] = 1;
			for (int j = 1; j <= m; ++j) if (vis[j]) ++ans;
			printf("%d\n", ans);
		}
		exit(0);
	}
}
namespace cpp2{
	stack<int>s;
	int dep[O], cnt[O], ans[O], Dep[O], son[O], S, T, res;
	il void Pop() {
		--cnt[c[s.top()]];
		if (!cnt[c[s.top()]]) --res;
		s.pop();
	}
	il void Push(int x) {
		s.push(x);
		if (!cnt[c[s.top()]]) ++res;
		++cnt[c[s.top()]];
	}
	il void dfs(int u, int fa) {
		if (!son[u]) return void(ans[u] = max(ans[u], res));
		int L = 0;
		for (int i = head[u]; i; i = e[i].nt)
			if (e[i].to != fa && e[i].to != son[u] && L < dep[e[i].to])
				L = dep[e[i].to];
		while (!s.empty() && Dep[s.top()] >= Dep[u] - L) Pop();
		Push(u);
		dfs(son[u], u);
		for (int i = head[u]; i; i = e[i].nt)
			if (e[i].to != fa && e[i].to != son[u]) {
				while (!s.empty() && Dep[s.top()] >= Dep[u] - dep[son[u]]) Pop();
				Push(u); dfs(e[i].to, u);
			}
		while (!s.empty() && Dep[s.top()] >= Dep[u] - dep[son[u]]) Pop();
		ans[u] = max(ans[u], res);
	}
	il void get(int u, int fa) {
		Dep[u] = Dep[fa] + 1;
		son[u] = dep[u] = 0;
		for (int i = head[u]; i; i = e[i].nt) {
			int v = e[i].to;
			if (v == fa) continue;
			get(v, u);
			if (dep[v] > dep[son[u]]) son[u] = v;
		}
		dep[u] = dep[son[u]] + 1;
	}
	il void main() {
		get(1, 0);
		for (int i = 1; i <= n; ++i) if (Dep[S] < Dep[i]) S = i;
		get(S, 0); dfs(S, 0);
		for (int i = 1; i <= n; ++i) if (Dep[T] < Dep[i]) T = i;
		get(T, 0); dfs(T, 0);
		for (int i = 1; i <= n; ++i) printf("%d\n", ans[i]);
	}
}
int main() {
	n = gi(), m = gi();
	for (int i = 1; i < n; ++i) {
		int u = gi(), v = gi();
		add(u, v); add(v, u);
	}
	for (int i = 1; i <= n; ++i) c[i] = gi();
	if (n <= 2000) cpp1::main();
	cpp2::main();
	return 0;
}