【JZOJ 杂题选讲】「JOI 2019 Final」独特的城市

题目

思路

每个点的答案一定在他到最远点的路径上。
每个点的最远点一定是直径的一个端点。

以直径的端点做根，每一个点x在距离较大的那个根上统计答案，就是求每个点x到根路径上的合法的y的不同颜色数。

进行长链剖分，md[x]表示x到子树里最远点的距离

假设现在在x上，用一个栈维护到x到根路径上剩下的合法的y。
如果往重儿子走，那么栈中到x距离<=(max(md[轻儿子])+1)的会被弹掉，
如果往轻儿子走，那么栈中到x距离<=(md[重儿子]+1)的会被弹掉。

由于元素加入次数是n-1，所以复杂度是O(n)。

在栈操作的同时维护一个桶即可统计栈内不同颜色数。

代码

#include<bits/stdc++.h>
using namespace std;
const int N=4e5+77;
int n,m,h[N],v[N * 2],nxt[N * 2],ec,dep[N],a,b,c[N],r[N],s[N],tp,cc[N],l[N],va,p[N];
void add(int x,int y)
{
	v[++ec] = y;
	nxt[ec] = h[x];
	h[x] = ec;
}
void dfs(int x,int fa)
{
	for(int i = h[x]; i; i = nxt[i])
		if(v[i] != fa)
		{
			dep[v[i]] = dep[x] + 1;
			dfs(v[i],x);
		}
}
void ins(int x)
{
	s[++tp] = x;
	cc[c[x]]++;
	if(cc[c[x]] == 1)
		va++;
}
void dl()
{
	cc[c[s[tp]]]--;
	if(!cc[c[s[tp]]])
		va--;
	tp--;
}
void d1(int x,int fa)
{
	dep[x] = dep[fa] + 1;
	p[x] = 0;
	for(int i = h[x]; i; i = nxt[i])
		if(v[i] != fa)
		{
			d1(v[i],x);
			if(l[v[i]] > l[p[x]])
				p[x] = v[i];
		}
	l[x] = l[p[x]] + 1;
}
void d2(int x,int fa)
{
	if(!p[x])
	{
		r[x] = max(r[x],va);
		return;
	}
	int m = 0;
	for(int i = h[x]; i; i = nxt[i])
		if(v[i] != fa && v[i] != p[x])
			m = max(m,l[v[i]]);
	while (tp && dep[s[tp]] >= dep[x] - m) dl();
	ins(x);
	d2(p[x],x);
	for(int i = h[x]; i; i = nxt[i])
		if(v[i] != fa && v[i] != p[x]) {
			while (tp && dep[s[tp]] >= dep[x] - l[p[x]]) dl();
			ins(x);
			d2(v[i],x);
		}
	while (tp && dep[s[tp]] >= dep[x] - l[p[x]]) dl();
	r[x] = max(r[x],va);
}
int main()
{
	cin >> n >> m;
	for(int i = 1; i < n; i++)
	{
		int x,y;
		cin >> x >> y;
		add(x,y);
		add(y,x);
	}
	for(int i = 1; i <= n; i++) cin >> c[i];
	dfs(1,0);
	for(int i = 1; i <= n; i++)
		if(dep[i] > dep[a])
			a = i;
	dfs(a,0);
	for(int i = 1; i <= n; i++)
		if(dep[i] > dep[b])
			b = i;
	d1(a,0);
	d2(a,0);
	d1(b,0);
	d2(b,0);
	for(int i = 1; i <= n; i++) printf("%d\n",r[i]);
}