洛谷P2664 樹上游戲（點分治）

傳送門

題解

　　由於一個sb錯誤調了一個晚上……鬼曉得我爲何$solve(rt)$會寫成$solve(v)$啊！！！一個$O(logn)$被我硬生生寫成$O(n)$了居然還能過$5$個點……話說還一直覺得只有動態點分會很難沒想到通常點分都這麼可啪……%%%大佬

　　咱們考慮一下，對於一棵樹，咱們要處理的是子樹對根的答案的貢獻，以及通過根的路徑的貢獻（也就是$LCA$爲根的點對的答案）。

　　對於樹中的一個點$i$，若是$i$的顏色是在$i$到根的路徑上第一次出現，那麼全部與$i$的$LCA$爲根的點，都能與$i$的子樹造成點對，且$i$的顏色會對那些點產生產生貢獻$size[i]$（先不考慮其餘子樹中是否有相同顏色）。

　　那麼咱們考慮$dfs$一遍整棵樹，並記錄，全部顏色產生的貢獻$col[i]$以及貢獻總和$sum$。而後$dfs$子樹，並把其中的全部顏色的貢獻給減掉。而後考慮子樹中的一個點，設$x$爲到根上的全部點的貢獻總和，$num$表示根到節點上的顏色個數，$y=size[root]-size[該子樹]$，那麼$ans[j]+=sum-x+num*size[y]$

　　而後最後記得加上對根節點的答案的貢獻$ans[root]+=sum-col[根的顏色]+size[root]$

　　不得不說思路真的挺妙的

  1 //minamoto
  2 #include<iostream>
  3 #include<cstdio>
  4 #include<cstring>
  5 #define ll long long
  6 using namespace std;
  7 #define getc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
  8 char buf[1<<21],*p1=buf,*p2=buf;
  9 template<class T>inline bool cmax(T&a,const T&b){return a<b?a=b,1:0;}
 10 inline int read(){
 11     #define num ch-'0'
 12     char ch;bool flag=0;int res;
 13     while(!isdigit(ch=getc()))
 14     (ch=='-')&&(flag=true);
 15     for(res=num;isdigit(ch=getc());res=res*10+num);
 16     (flag)&&(res=-res);
 17     #undef num
 18     return res;
 19 }
 20 char sr[1<<21],z[20];int C=-1,Z;
 21 inline void Ot(){fwrite(sr,1,C+1,stdout),C=-1;}
 22 inline void print(ll x){
 23     if(C>1<<20)Ot();if(x<0)sr[++C]=45,x=-x;
 24     while(z[++Z]=x%10+48,x/=10);
 25     while(sr[++C]=z[Z],--Z);sr[++C]='\n';
 26 }
 27 const int N=200005;
 28 int head[N],Next[N<<1],ver[N<<1],c[N],son[N],sz[N],size,cnt[N];
 29 ll col[N],ans[N],much,sum,num;
 30 int tot,n,rt;bool vis[N];
 31 inline void add(int u,int v){
 32     ver[++tot]=v,Next[tot]=head[u],head[u]=tot;
 33     ver[++tot]=u,Next[tot]=head[v],head[v]=tot;
 34 }
 35 void findrt(int u,int fa){
 36     sz[u]=1,son[u]=0;
 37     for(int i=head[u];i;i=Next[i]){
 38         int v=ver[i];
 39         if(!vis[v]&&v!=fa){
 40             findrt(v,u);
 41             sz[u]+=sz[v];
 42             cmax(son[u],sz[v]);
 43         }
 44     }
 45     cmax(son[u],size-sz[u]);
 46     if(son[u]<son[rt]) rt=u;
 47 }
 48 void dfs1(int u,int fa){
 49     //要從新dfs一次，不能直接在找根時的樹上作（否則子樹和之類的會出錯） 
 50     //順便維護各類東西 
 51     sz[u]=1,++cnt[c[u]];
 52     for(int i=head[u];i;i=Next[i]){
 53         int v=ver[i];
 54         if(!vis[v]&&v!=fa){
 55             dfs1(v,u),sz[u]+=sz[v];
 56         }
 57     }
 58     if(cnt[c[u]]==1) sum+=sz[u],col[c[u]]+=sz[u];
 59     --cnt[c[u]];
 60 }
 61 void change(int u,int fa,int val){
 62     ++cnt[c[u]];
 63     for(int i=head[u];i;i=Next[i]){
 64         int v=ver[i];
 65         if(!vis[v]&&v!=fa) change(v,u,val);
 66     }
 67     if(cnt[c[u]]==1) sum+=sz[u]*val,col[c[u]]+=sz[u]*val;
 68     --cnt[c[u]];
 69 }
 70 void dfs2(int u,int fa){
 71     //把這棵子樹裏的顏色的影響消除掉
 72     //順便更新答案 
 73     ++cnt[c[u]];
 74     if(cnt[c[u]]==1) sum-=col[c[u]],++num;
 75     ans[u]+=sum+num*much;
 76     for(int i=head[u];i;i=Next[i]){
 77         int v=ver[i];
 78         if(!vis[v]&&v!=fa) dfs2(v,u);
 79     }
 80     if(cnt[c[u]]==1) sum+=col[c[u]],--num;
 81     --cnt[c[u]];
 82 }
 83 void clear(int u,int fa){
 84     cnt[c[u]]=col[c[u]]=0;
 85     for(int i=head[u];i;i=Next[i]){
 86         int v=ver[i];
 87         if(!vis[v]&&v!=fa) clear(v,u);
 88     }
 89 }
 90 void did(int u){
 91     //直接帶進去亂搞 
 92     dfs1(u,0);
 93     ans[u]+=sum-col[c[u]]+sz[u];
 94     for(int i=head[u];i;i=Next[i]){
 95         int v=ver[i];
 96         if(!vis[v]){
 97             //dfs，而後把各類影響消除掉 
 98             ++cnt[c[u]];
 99             sum-=sz[v];
100             col[c[u]]-=sz[v];
101             change(v,u,-1);
102             --cnt[c[u]];
103             much=sz[u]-sz[v];
104             dfs2(v,u);
105             ++cnt[c[u]];
106             sum+=sz[v];
107             col[c[u]]+=sz[v];
108             change(v,u,1);
109             --cnt[c[u]];
110         }
111     }
112     sum=0,num=0,clear(u,0);
113 }
114 void solve(int u){
115     did(u),vis[u]=true;
116     for(int i=head[u];i;i=Next[i]){
117         int v=ver[i];
118         if(!vis[v]){
119             rt=0,size=sz[v];
120             findrt(v,0),solve(rt);
121         }
122     }
123 }
124 int main(){
125     n=read();
126     for(int i=1;i<=n;++i) c[i]=read();
127     for(int i=1;i<n;++i){
128         int u=read(),v=read();
129         add(u,v);
130     }
131     son[0]=n+1,rt=0,size=n;
132     findrt(1,0),solve(rt);
133     for(int i=1;i<=n;++i) print(ans[i]);
134     Ot();
135     return 0;
136 }