题目：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; }
原文：https://www.cnblogs.com/Narh/p/10608298.html