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

时间：2022-10-27 14:59:23

题目：

　　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+树上统计

相关文章