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+樹上統計