题目：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 ",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) ",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 ",ans[i]); return 0; }
看了题解，还有一种 O(n) 的做法。https://www.luogu.org/problemnew/solution/P2664
只看某种颜色的点，删掉它们，树变成一些 “小树” 。
该颜色对一个不是该颜色的点的贡献就是 n - siz[ 该点所在小树 ] 。
所以考虑在遍历树的时候把这样的 siz[ 小树 ] 都找出来。考虑一个某颜色下小树的 siz 记录在该小树最浅的节点上。这样 “该颜色” 一定是该点的父亲的颜色，且该点和父亲不同色。所以每个点只会记录一个信息。也就是想对每个点求出这部分的大小：
（其中括住的那些点都是和父亲一个颜色）考虑令 sc[ i ] 表示该点子树里最浅点颜色为 i 的这种块的大小，则一个点记录的信息 f[ cr ] = siz[ cr ] - sc [ c[ fa ] ] （siz 是子树大小）。sc[ i ] 其实可以全局开一个，只要记录一下进 cr 子树之前的 sc 值，和遍历完 cr 子树的 sc 值对比一下就能得到子树里的 sc 值了。考虑统计答案的时候，如何知道这个点所在的其他颜色（非父亲颜色）下的“小树”大小？
其实这就是该点所有祖先的 f[ ] ，“父亲颜色” 相同的只保留最深的。这个只要开一个对于颜色的全局数组记录目前各颜色用的 f[ ] 值是多少就行了。
#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 Mx(int a,int b){return a>b?a:b;} const int N=1e5+5; int n,c[N],hd[N],xnt,to[N<<1],nxt[N<<1]; int siz[N],sc[N],f[N],ct; ll ans[N],sm; bool vis[N]; void add(int x,int y){to[++xnt]=y;nxt[xnt]=hd[x];hd[x]=xnt;} void dfs(int cr,int fa) { siz[cr]=1; int nw=c[fa], tp=sc[nw]; for(int i=hd[cr],v;i;i=nxt[i]) if((v=to[i])!=fa) dfs(v,cr), siz[cr]+=siz[v]; sc[c[cr]]++; if(!fa)return; f[cr]=siz[cr]-(sc[nw]-tp); sc[nw]=tp+siz[cr]; } void dfsx(int cr,int fa) { int nw=c[fa], tp=sc[nw]; if(fa){ sm-=sc[nw]; sc[nw]=f[cr]; sm+=sc[nw];} ans[cr]=(ll)n*ct-sm+sc[c[cr]]; for(int i=hd[cr],v;i;i=nxt[i]) if((v=to[i])!=fa) dfsx(v,cr); if(fa){ sm-=sc[nw]; sc[nw]=tp; sm+=sc[nw];} } int main() { n=rdn(); int mx=0; for(int i=1;i<=n;i++) { c[i]=rdn();mx=Mx(mx,c[i]); if(!vis[c[i]])vis[c[i]]=1, ct++; } for(int i=1,u,v;i<n;i++) {u=rdn();v=rdn();add(u,v);add(v,u);} dfs(1,0); for(int i=1;i<=mx;i++)if(vis[i])sc[i]=n-sc[i], sm+=sc[i];//vis!!! dfsx(1,0); for(int i=1;i<=n;i++)printf("%lld ",ans[i]); return 0; }