題目：https://www.luogu.org/problemnew/show/P2664

想到點分治的話，想到一條路徑會對路徑的兩端點的 ans[ ] 有貢獻，所以每次就考慮過重心的路徑。

先想的是枚舉顏色，每次做一種顏色。

枚舉重心的一個子樹的時候，就記錄一下當前重心 cr 的其它子樹裏有幾條 “從 cr 開始，含有該顏色” 的路徑，設為 sc；

在當前路徑上如果沒碰見有顏色的點，那麽給當前點 ans[ ] += sc ；如果碰見過有顏色的點，就給當前點 ans[ ] += siz(其他子樹)+1 （+1表示重心那個點）；

先把所有子樹的 sc 累加起來，進入這棵子樹之前原樣 dfs 一下該子樹把它帶來的 sc 減去，再進入它算答案即可。

然後考慮所有顏色也可以一起做。就是把 sc 擴充成一個數組，用 vc[ i ] 表示當前路徑有沒有曾經有過 i 這種顏色（若使 vc = 1 的點是 cr ，則 use[ cr ] = 1 ，方便回溯的時候把 vc 改回來）。

先把所有顏色的 sc（其他子樹） 累加起來，記為 s1 ；如果當前路徑碰到一種顏色 i ，就在 s1 裏減掉 sc[ i ]（其他子樹），而給一個初始為 0 的 lj += siz（其他子樹），表示該顏色的貢獻不要求 “其他子樹路徑有該顏色” 了。

註意重心那個點也對 vc 和 sc 有影響。並註意最後刪除影響的時候記得刪除重心帶來的影響！

#include<cstdio>
#include<cstring>
#include<algorithm>
#define ll long long
using namespace std;
int rdn()
{
  int ret=0;bool fx=1;char ch=getchar();
  while(ch>‘9‘||ch<‘0‘){if(ch==‘-‘)fx=0;ch=getchar();}
  while(ch>=‘0‘&&ch<=‘9‘)ret=ret*10+ch-‘0‘,ch=getchar();
  
 
return fx?ret:-ret;
}
int Mn(int a,int b){return a<b?a:b;}
int Mx(int a,int b){return a>b?a:b;}
const int N=1e5+5,mod=1e9+7;
int n,c[N],siz[N],hd[N],xnt,to[N<<1],nxt[N<<1],rt,mn;
ll ans[N],sc[N],s1,s2; bool vis[N],use[N],vc[N];
void add(int x,int y){to[++xnt]=y;nxt[xnt]=hd[x];hd[x]=xnt;}
void get_rt(int cr,int fa,int s)
{
  siz[cr]=1; int mx=0;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa)
      {
    get_rt(v,cr,s); siz[cr]+=siz[v];
    mx=Mx(mx,siz[v]);
      }
  mx=Mx(mx,s-siz[cr]); if(mx<mn)mn=mx,rt=cr;
}
void ini_dfs(int cr,int fa)
{
  siz[cr]=1;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa)
      ini_dfs(v,cr), siz[cr]+=siz[v];
}
void dfs(int cr,int fa,int fx)
{
  int nw=c[cr];
  if(!vc[nw])
    {
      vc[nw]=1; use[cr]=1;
      int d=(fx?siz[cr]:-siz[cr]); sc[nw]+=d; s1+=d;
      if(fx==1)ans[rt]+=siz[cr];
    }
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa) dfs(v,cr,fx);
  if(use[cr]){vc[nw]=0; use[cr]=0;}
}
void dfsx(int cr,int fa,ll lj)
{
  int nw=c[cr];
  if(!vc[nw])
    { vc[nw]=1; use[cr]=1; s1-=sc[nw]; lj+=s2;}
  //printf(" dfsx:cr=%d s1=%lld lj=%lld\n",cr,s1,lj);
  ans[cr]+=s1+lj;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa) dfsx(v,cr,lj);
  if(use[cr]){vc[nw]=0; use[cr]=0; s1+=sc[nw];}
}
void del_dfs(int cr,int fa)
{
  sc[c[cr]]=0;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]]&&v!=fa) del_dfs(v,cr);
}
void solve(int cr,int s)
{
  //printf("solve:cr=%d (s=%d)\n",cr,s);
  vis[cr]=1; int nw=c[cr];
  vc[nw]=1; sc[nw]=s1=s; ans[rt]+=s;//s not s-1 for self
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])ini_dfs(v,cr),dfs(v,cr,1);
  s2=s;
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])
      {
    dfs(v,cr,0); sc[nw]-=siz[v]; s1-=siz[v]; s2=s-siz[v];
    /*printf(" v=%d(s1=%lld s2=%lld):",v,s1,s2);
      for(int j=1;j<=2;j++)printf("%lld ",sc[j]);puts("");*/
    dfsx(v,cr,0); sc[nw]+=siz[v]; s1+=siz[v];
    dfs(v,cr,2);//
      }
  /*printf("ans:");
    for(int i=1;i<=n;i++)printf("%lld ",ans[i]);puts("");*/
  vc[nw]=sc[nw]=0;///!!!
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])del_dfs(v,cr);
  for(int i=hd[cr],v;i;i=nxt[i])
    if(!vis[v=to[i]])
      {
    mn=N;get_rt(v,0,siz[v]);solve(rt,siz[v]);
      }
}
int main()
{
  n=rdn();
  for(int i=1;i<=n;i++)c[i]=rdn();
  for(int i=1,u,v;i<n;i++)
    { u=rdn();v=rdn();add(u,v);add(v,u);}
  mn=N;get_rt(1,0,n);solve(rt,n);
  for(int i=1;i<=n;i++)printf("%lld\n",ans[i]);
  return 0;
}

洛谷 2664 樹上遊戲——點分治