题目
题解
首先发现可能能计入贡献的点都在直径上
所以先$dfs$两遍把直径搞出来,然后从直径两端分别$dfs$,只统计祖先部分(链上部分)是否有贡献,用子树中的最深点把距离重复的祖先去除掉,发现这个东西可以用单调栈维护
两遍求答案取最大值即可
代码
# include<iostream>
# include<cstring>
# include<cstdio>
# include<algorithm>
using namespace std;
const int MAX=2e5+5;
struct p{
int x,y;
}c[MAX<<1];
int n,m,num,Top,ans,s1,s2,maxn;
int h[MAX],col[MAX],st[MAX],d[MAX],D[MAX],son[MAX],Num[MAX],Ans[MAX];
void add(int x,int y)
{
c[++num]=(p){h[x],y},h[x]=num;
c[++num]=(p){h[y],x},h[y]=num;
}
int read()
{
int x(0);
char ch=getchar();
for(;!isdigit(ch);ch=getchar());
for(;isdigit(ch);x=x*10+ch-48,ch=getchar());
return x;
}
void add(int x) {if(++Num[col[x]]==1) ++ans;}
void del(int x) {if(--Num[col[x]]==0) --ans;}
void Add(int x) {add(st[++Top]=x);}
void Del() {del(st[Top--]);}
void dfs(int x,int &s,int fa=0)
{
if((d[x]=d[fa]+1)>maxn) maxn=d[x],s=x;
for(int i=h[x];i;i=c[i].x)
if(c[i].y^fa) dfs(c[i].y,s,x);
}
void dfs_(int x,int fa=0)
{
d[x]=d[fa]+1,D[x]=0,D[son[x]=0]=-1;
for(int i=h[x];i;i=c[i].x)
if(c[i].y^fa)
{
dfs_(c[i].y,x),D[x]=max(D[x],D[c[i].y]+1);
if(D[c[i].y]>D[son[x]]) son[x]=c[i].y;
}
}
void dfs__(int x,int fa=0)
{
if(!son[x]) return void(Ans[x]=max(Ans[x],ans));
int mx=0;
for(int i=h[x];i;i=c[i].x)
if((c[i].y^fa)&&(c[i].y^son[x])) mx=max(mx,D[c[i].y]+1);
while(Top&&d[st[Top]]+mx>=d[x]) Del();
Add(x),dfs__(son[x],x);
for(int i=h[x];i;i=c[i].x)
if((c[i].y^fa)&&(c[i].y^son[x]))
{
while(Top&&d[st[Top]]+D[son[x]]+1>=d[x]) Del();
Add(x),dfs__(c[i].y,x);
}
while(Top&&d[st[Top]]+D[son[x]]+1>=d[x]) Del();
Ans[x]=max(Ans[x],ans);
}
int main()
{
n=read(),m=read();
for(int i=1;i<n;++i)
add(read(),read());
for(int i=1;i<=n;++i)
col[i]=read();
dfs(1,s1),maxn=0,dfs(s1,s2),dfs_(s1),dfs__(s1),dfs_(s2),dfs__(s2);
for(int i=1;i<=n;++i)
printf("%d\n",Ans[i]);
return 0;
}