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

题目：

　　P2664 树上游戏

 

分析：

　　原本是练习点分治的时候看到了这道题。无心中发现题解中有一种方法能够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总和 

dfs+树上统计