Luogu P2664 树上游戏 dfs+树上统计
题目：
分析：
本来是练习点分治的时候看到了这道题。无意中发现题解中有一种方法可以O(N)解决这道题，就去膜拜了一下。
这个方法是，假如对于某一种颜色，将所有这种颜色的点全部删去，原树会被割成若干棵小树，那么这个颜色对每个点的贡献就是：树的大小n - 所在小树的大小sz。所以我们要求出对于每个点来说，删去所有这个点颜色的点，这个点以下的小树size，这个用一个dfs和一个类似前缀和相减的过程，就可以求出。
接下来统计每个点的答案，就是所有颜色对这个点的贡献：n*颜色数-对于每种颜色这个点所处的小树size，这个用一遍dfs也可以求出来，类似容斥的思想。
代码中有注释：
代码：
1 #include<bits/stdc++.h> 2 #define ll long long 3 using namespace std; 4 const int N=100005; 5 struct node{int y,nxt;}e[N*2]; 6 ll n,sm,qwq,vis[N],ans[N],c=0,h[N]; 7 ll col[N],siz[N],tmp[N],lz[N],sn[N]; 8 void add(int x,int y){ 9 e[++c]=(node){y,h[x]};h[x]=c; 10 e[++c]=(node){x,h[y]};h[y]=c; 11 } void dfs(int x,int fa){ 12 siz[x]=1;ll cnt=tmp[col[fa]]; 13 for(int i=h[x],y;i;i=e[i].nxt) 14 if((y=e[i].y)!=fa) dfs(y,x),siz[x]+=siz[y]; 15 tmp[col[x]]++;if(fa){ 16 lz[x]=siz[x]-tmp[col[fa]]+cnt; 17 tmp[col[fa]]+=lz[x]; 18 } return ; 19 } void get(int x,int fa){ 20 ll sgn=sn[col[fa]]; 21 qwq+=lz[x]-sn[col[fa]]; 22 sn[col[fa]]=lz[x]; 23 ans[x]=n*sm-qwq+sn[col[x]]; 24 for(int i=h[x],y;i;i=e[i].nxt) 25 if((y=e[i].y)!=fa) get(y,x); 26 sn[col[fa]]=sgn; 27 qwq-=lz[x]-sn[col[fa]]; 28 } int main(){ 29 scanf("%lld",&n); 30 for(int i=1;i<=n;i++){ 31 scanf("%lld",&col[i]); 32 if(!vis[col[i]]) 33 vis[col[i]]=1,sm++; 34 } for(int i=1,x,y;i<n;i++) 35 scanf("%d%d",&x,&y),add(x,y); 36 dfs(1,0); 37 for(int i=1;i<=100000;i++) 38 if(vis[i]) qwq+=n-tmp[i], 39 sn[i]=n-tmp[i];get(1,0); 40 for(int i=1;i<=n;i++) 41 printf("%lld\n",ans[i]);return 0; 42 }//lz[x]表示将x的父节点这种颜色的点全部删掉余下的这个 43 //小树的size 44 //sn记录的是对于每种颜色，当前深度以上的最近的小树sz 45 //qwq变量计算的是当前位置对于每种颜色所处的小树sz总和
