首先，定义“点 $t$ 对点 $u$ 有贡献”等价于 $t$ 是 $u$ 的独特的城市。

钦定点 $u$ 为根，令点集 $L_u$ 表示到所有到点 $u$ 距离最远的叶子节点构成的集合，对点 $u$ 有贡献的点一定在 $L_u$ 中所有点到 $u$ 的路径的交上。

注意，路径交上的所有点并不都一定对 $u$ 有贡献。

因此一个想法就是考虑找到随便一个点 $v$ 距离点 $u$ 最远，考察 $u,v$ 两点间的路径上的点。

而考虑离点 $u$ 最远的点，可以选一条直径 $s-t$，$s,t$ 两个点中一定存在一个点是距离点 $u$ 最远的点。

进而问题就转化成了：给定一个点 $rt$，以 $rt$ 为根，求出每个点 $u$ 到 $rt$ 路径上所有对点 $u$ 有贡献的点的颜色数。

分别以 $s,t$ 为 $rt$ 求一遍上面的问题即可。

接下来考虑如何去求解上面转化后的问题。

考虑一个在点 $u$ 到根节点路径上的点 $p$ 不会对点 $u$ 产生贡献的充分必要条件：存在一个点 $t$ 使得 $t$ 与 $u$ 的距离等于 $p$ 与 $u$ 的距离，按照点 $t$ 的位置分类讨论：

• $t$ 在 $u$ 子树内。

• $t$ 在 $u$ 子树外。

考虑维护出一个点集 $T_u$，$\forall x\in T_u$ 满足：

• $x$ 在根节点到 $fa[u]$ 的路径上。

• 不存在一个点 $t$ 在 $u$ 的子树外，使得 $dist(u,t)=dist(u,x)$。

同时考虑点集 $S_u$，$\forall x\in S_u$ 满足：

• $x$ 在根节点到 $u$ 的路径上。

• 不存在一个点 $t$ ，使得 $dist(u,t)=dist(u,x)$。

并且 $S_u\subseteq T_u$，删掉 $T_u$ 中所有满足存在一个点 $t$ 在 $u$ 子树内，使得 $dist(u,x)=dist(u,t)$ 的点 $x$ 即可将 $T_u$ “变成” $S_u$。

考虑 $u$ 的一个儿子 $v$，想要构造出 $T_v$，需要对 $T_u$ 进行哪些 “改造”（或者说点集如何变化）。

$v$ 的子树外等价于 $fa[v]$ 的子树外并上 $v$ 的兄弟子树

• 先令 $T_v\gets T_u\cup \{u\}$。

• 删掉 $T_v$ 中所有满足存在一个点 $t$ 满足 $t$ 属于 $v$ 的兄弟的子树中，$dist(u,t) = dist(u,x)$ 的点 $x$。

而 “删掉 $T_v$ 中所有满足存在一个点 $t$ 满足 $t$ 属于 $v$ 的兄弟的子树中，$dist(u,t) = dist(u,x)$ 的点 $x$”，可以转化成 “删掉 $T_v$ 中所有 $dist(u,x) <= \text{v 的所有兄弟子树的最大深度 + 1}$ 的点 $x$。

考虑对该树做长链剖分，令点 $u$ 的长儿子为 $son[u]$，$\mathrm{maxdist(u)}$ 表示 $u$ 子树内距离 $u$ 距离最远的点与 $u$ 的距离，$ts[u]$ 表示 $u$ 的次长儿子，$\mathrm{len(u)}=\mathrm{maxdist(ts[u])}+1$。

注意到，对于 $u$ 的所有非长儿子，在点集 $T_u\cup\{u\}$ 中要删除的点集都一样。

而对于点 $u$ 的长儿子，要删除的点集为 $T_u\cup\{u\}$ 中所有满足 $dist(u,x)\le \mathrm{len(u)}$ 的点 $x$。

考虑在 DFS 的过程中维护点集 $T_u$，在进入子树 $u$ 前，全局维护的信息应为 $T_u$ 的信息。

进入子树 $u$ 后依次进行以下操作：

• 删除当前全局维护的点集中所有满足 $dist(u,x)\le \mathrm{len(u)}$ 的点 $x$。

• 在当前全局维护的点集中加入点 $u$。

• 递归长子树。

• 删掉当前全局维护的点集中所有满足 $dist(u,x)\le \mathrm{maxdist(u)}$，当前全局维护的点集为 $S_u$，同时也是对于任意轻儿子 $v$ 的 $T_v$。

• 依次递归轻子树,其中在递归轻子树前需要将点 $u$ 加入全局维护的点集中。

• 如果点 $u$ 在全局维护的点集中，删除点 $u$。

这样操作的合法性：

在递归长子树时，在长子树内进行的 “删除操作”影响到的 $u$ 的祖先，至多影响到 $u$ 的 $\mathrm{maxdist(u)}-2$ 级祖先，而这些点在处理轻子树之前本来就需要被删除。

而在递归非长子树 $v$ 时，只会至多影响 $u$ 的 $\mathrm{maxdist(v)}-1$ 级祖先，而这些祖先一定都被删光光了，因此在处理轻子树时一定不会影响到集合中 $u$ 的祖先。

这样，一个元素至多贡献儿子个数次总复杂的为 $O(n)$。

D57 树的直径 树形DP+栈 P6118 [JOI 2019 Final] 独特的城市

// 树的直径 树形DP+栈 O(n)
#include<bits/stdc++.h>
using namespace std;

const int N=200010;
int n,m,c[N];
int h[N],to[N<<1],ne[N<<1],idx;
void adde(int x,int y){
  to[++idx]=y,ne[idx]=h[x],h[x]=idx;
}

int d[N],p,l,r;
void dfs1(int x,int fa){
  if(d[x]>d[p]) p=x;
  for(int i=h[x];i;i=ne[i]){
    int y=to[i];
    if(y!=fa){
      d[y]=d[x]+1; //记录从根到y的距离
      dfs1(y,x);
    }
  }
}
int dep[N],d1[N],d2[N],son[N];
void dfs2(int x,int fa){ //树形DP
  dep[x]=dep[fa]+1,d1[x]=d2[x]=0,son[x]=0;
  for(int i=h[x];i;i=ne[i]){
    int y=to[i];
    if(y!=fa){
      dfs2(y,x);
      if(d1[x]<d1[y]+1) d2[x]=d1[x],d1[x]=d1[y]+1,son[x]=y;
      else if(d2[x]<d1[y]+1) d2[x]=d1[y]+1;
    }
  }
}
int stk[N],top,num[N],res,ans[N];
void add(int x){++num[c[x]]; if((num[c[x]])==1)res++;} //加入点的贡献
void del(int x){--num[c[x]]; if((num[c[x]])==0)res--;} //删除点的贡献
void dfs3(int x,int fa){
  if(fa) add(stk[++top]=fa); //加入父节点
  while(top && dep[x]-dep[stk[top]]<=d2[x]) del(stk[top--]); //删除上面长度≤次长链的节点
  if(son[x]) dfs3(son[x],x); //先遍历长儿子
  while(top && dep[x]-dep[stk[top]]<=d1[x]) del(stk[top--]); //删除上面长度≤最长链的节点
  for(int i=h[x];i;i=ne[i]){ //后遍历短儿子
    int y=to[i];
    if(y!=fa&&y!=son[x]) dfs3(y,x);
  }
  ans[x]=max(ans[x],res); //更新答案
  if(stk[top]==fa) del(stk[top--]); //删除父节点，恢复现场
}
int main(){
  cin>>n>>m;
  for(int i=1,x,y;i<n;i++){
    cin>>x>>y;
    adde(x,y),adde(y,x);
  }
  for(int i=1;i<=n;i++) cin>>c[i];
  
  dfs1(1,0); l=p; d[p]=0;
  dfs1(p,0); r=p; //记录直径端点
  
  dfs2(l,0); //左端进入，记录深度、最长链、次长链、长儿子
  dfs3(l,0); //计算答案
  memset(num,0,sizeof(num)); //清空桶
  top=res=0; //清空栈
  dfs2(r,0); //右端进入，记录深度、最长链、次长链、长儿子
  dfs3(r,0); //计算答案
  for(int i=1;i<=n;i++) cout<<ans[i]<<"\n";
}