Luogu P2664 树上游戏 dfs+树上统计

题目：

　　P2664 树上游戏

分析：

　　本来是练习点分治的时候看到了这道题。无意中发现题解中有一种方法可以O(N)解决这道题，就去膜拜了一下。

　　这个方法是，假如对于某一种颜色，将所有这种颜色的点全部删去，原树会被割成若干棵小树，那么这个颜色对每个点的贡献就是：树的大小n - 所在小树的大小sz。所以我们要求出对于每个点来说，删去所有这个点颜色的点，这个点以下的小树size，这个用一个dfs和一个类似前缀和相减的过程，就可以求出。

　　接下来统计每个点的答案，就是所有颜色对这个点的贡献：n*颜色数-对于每种颜色这个点所处的小树size，这个用一遍dfs也可以求出来，类似容斥的思想。

　　代码中有注释：

代码：

 #include<bits/stdc++.h>
 #define ll long long
 using namespace std;
 const int N=;
 struct node{int y,nxt;}e[N*];
 ll n,sm,qwq,vis[N],ans[N],c=,h[N];
 ll col[N],siz[N],tmp[N],lz[N],sn[N];
 void add(int x,int y){
     e[++c]=(node){y,h[x]};h[x]=c;
     e[++c]=(node){x,h[y]};h[y]=c;
 } void dfs(int x,int fa){
     siz[x]=;ll cnt=tmp[col[fa]];
     for(int i=h[x],y;i;i=e[i].nxt)
     if((y=e[i].y)!=fa) dfs(y,x),siz[x]+=siz[y];
     tmp[col[x]]++;if(fa){
         lz[x]=siz[x]-tmp[col[fa]]+cnt;
         tmp[col[fa]]+=lz[x];
     } return ;
 } void get(int x,int fa){
     ll sgn=sn[col[fa]];
     qwq+=lz[x]-sn[col[fa]];
     sn[col[fa]]=lz[x];
     ans[x]=n*sm-qwq+sn[col[x]];
     for(int i=h[x],y;i;i=e[i].nxt)
     if((y=e[i].y)!=fa) get(y,x);
     sn[col[fa]]=sgn;
     qwq-=lz[x]-sn[col[fa]];
 } int main(){
     scanf("%lld",&n);
     for(int i=;i<=n;i++){
         scanf("%lld",&col[i]);
         if(!vis[col[i]])
         vis[col[i]]=,sm++;
     } for(int i=,x,y;i<n;i++)
     scanf("%d%d",&x,&y),add(x,y);
     dfs(,);
     for(int i=;i<=;i++)
     if(vis[i]) qwq+=n-tmp[i],
     sn[i]=n-tmp[i];get(,);
     for(int i=;i<=n;i++)
     printf("%lld\n",ans[i]);return ;
 }//lz[x]表示将x的父节点这种颜色的点全部删掉余下的这个
 //小树的size
 //sn记录的是对于每种颜色，当前深度以上的最近的小树sz
 //qwq变量计算的是当前位置对于每种颜色所处的小树sz总和 

dfs+树上统计