洛谷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]$

　　不得不说思路真的挺妙的

 //minamoto
 #include<iostream>
 #include<cstdio>
 #include<cstring>
 #define ll long long
 using namespace std;
 #define getc() (p1==p2&&(p2=(p1=buf)+fread(buf,1,1<<21,stdin),p1==p2)?EOF:*p1++)
 char buf[<<],*p1=buf,*p2=buf;
 template<class T>inline bool cmax(T&a,const T&b){return a<b?a=b,:;}
 inline int read(){
     #define num ch-'0'
     char ch;bool flag=;int res;
     while(!isdigit(ch=getc()))
     (ch=='-')&&(flag=true);
     for(res=num;isdigit(ch=getc());res=res*+num);
     (flag)&&(res=-res);
     #undef num
     return res;
 }
 char sr[<<],z[];int C=-,Z;
 inline void Ot(){fwrite(sr,,C+,stdout),C=-;}
 inline void print(ll x){
     if(C><<)Ot();if(x<)sr[++C]=,x=-x;
     while(z[++Z]=x%+,x/=);
     while(sr[++C]=z[Z],--Z);sr[++C]='\n';
 }
 const int N=;
 int head[N],Next[N<<],ver[N<<],c[N],son[N],sz[N],size,cnt[N];
 ll col[N],ans[N],much,sum,num;
 int tot,n,rt;bool vis[N];
 inline void add(int u,int v){
     ver[++tot]=v,Next[tot]=head[u],head[u]=tot;
     ver[++tot]=u,Next[tot]=head[v],head[v]=tot;
 }
 void findrt(int u,int fa){
     sz[u]=,son[u]=;
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]&&v!=fa){
             findrt(v,u);
             sz[u]+=sz[v];
             cmax(son[u],sz[v]);
         }
     }
     cmax(son[u],size-sz[u]);
     if(son[u]<son[rt]) rt=u;
 }
 void dfs1(int u,int fa){
     //要重新dfs一次，不能直接在找根时的树上做（不然子树和之类的会出错）
     //顺便维护各种东西
     sz[u]=,++cnt[c[u]];
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]&&v!=fa){
             dfs1(v,u),sz[u]+=sz[v];
         }
     }
     if(cnt[c[u]]==) sum+=sz[u],col[c[u]]+=sz[u];
     --cnt[c[u]];
 }
 void change(int u,int fa,int val){
     ++cnt[c[u]];
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]&&v!=fa) change(v,u,val);
     }
     if(cnt[c[u]]==) sum+=sz[u]*val,col[c[u]]+=sz[u]*val;
     --cnt[c[u]];
 }
 void dfs2(int u,int fa){
     //把这棵子树里的颜色的影响消除掉
     //顺便更新答案
     ++cnt[c[u]];
     if(cnt[c[u]]==) sum-=col[c[u]],++num;
     ans[u]+=sum+num*much;
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]&&v!=fa) dfs2(v,u);
     }
     if(cnt[c[u]]==) sum+=col[c[u]],--num;
     --cnt[c[u]];
 }
 void clear(int u,int fa){
     cnt[c[u]]=col[c[u]]=;
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]&&v!=fa) clear(v,u);
     }
 }
 void did(int u){
     //直接带进去乱搞
     dfs1(u,);
     ans[u]+=sum-col[c[u]]+sz[u];
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]){
             //dfs，然后把各种影响消除掉
             ++cnt[c[u]];
             sum-=sz[v];
             col[c[u]]-=sz[v];
             change(v,u,-);
             --cnt[c[u]];
             much=sz[u]-sz[v];
             dfs2(v,u);
             ++cnt[c[u]];
             sum+=sz[v];
             col[c[u]]+=sz[v];
             change(v,u,);
             --cnt[c[u]];
         }
     }
     sum=,num=,clear(u,);
 }
 void solve(int u){
     did(u),vis[u]=true;
     for(int i=head[u];i;i=Next[i]){
         int v=ver[i];
         if(!vis[v]){
             rt=,size=sz[v];
             findrt(v,),solve(rt);
         }
     }
 }
 int main(){
     n=read();
     for(int i=;i<=n;++i) c[i]=read();
     for(int i=;i<n;++i){
         int u=read(),v=read();
         add(u,v);
     }
     son[]=n+,rt=,size=n;
     findrt(,),solve(rt);
     for(int i=;i<=n;++i) print(ans[i]);
     Ot();
     return ;
 }