洛谷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[1<<21],*p1=buf,*p2=buf;
template<class T>inline bool cmax(T&a,const T&b){return a<b?a=b,1:0;}
inline int read(){
    #define num ch-'0'
    char ch;bool flag=0;int res;
    while(!isdigit(ch=getc()))
    (ch=='-')&&(flag=true);
    for(res=num;isdigit(ch=getc());res=res*10+num);
    (flag)&&(res=-res);
    #undef num
    return res;
}
char sr[1<<21],z[20];int C=-1,Z;
inline void Ot(){fwrite(sr,1,C+1,stdout),C=-1;}
inline void print(ll x){
    if(C>1<<20)Ot();if(x<0)sr[++C]=45,x=-x;
    while(z[++Z]=x%10+48,x/=10);
    while(sr[++C]=z[Z],--Z);sr[++C]='\n';
}
const int N=200005;
int head[N],Next[N<<1],ver[N<<1],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]=1,son[u]=0;
    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]=1,++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]]==1) 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]]==1) 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]]==1) 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]]==1) sum+=col[c[u]],--num;
    --cnt[c[u]];
}
void clear(int u,int fa){
    cnt[c[u]]=col[c[u]]=0;
    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,0);
    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,-1);
            --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,1);
            --cnt[c[u]];
        }
    }
    sum=0,num=0,clear(u,0);
}
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=0,size=sz[v];
            findrt(v,0),solve(rt);
        }
    }
}
int main(){
    n=read();
    for(int i=1;i<=n;++i) c[i]=read();
    for(int i=1;i<n;++i){
        int u=read(),v=read();
        add(u,v);
    }
    son[0]=n+1,rt=0,size=n;
    findrt(1,0),solve(rt);
    for(int i=1;i<=n;++i) print(ans[i]);
    Ot();
    return 0;
}