洛谷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
  3 #include
  4 #include
  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 a1: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]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;ii){
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 }