点分治-luogu2664

思路

这题是个很好的点分治题目。
首先明确我们的目的，既然使用了点分治，那么就需要O(n)的时间内处理出重心rt以及“rt的儿子”的答案。
考虑重心的一个子树的根节点，如果这个点v的颜色是第一次出现，那么它对其他不在该子树的点的答案的贡献(包括rt)就是size[v]
然后我们再考虑这个子树，注意刚刚我们考虑的是子树的根节点，现在考虑的是整个子树对其他“rt儿子”的贡献。假设该子树的根节点是v，从v~rt上有num种颜色，那么对其他节点的贡献就是num*(size[rt]-size[v])
整体的步骤：

1. dfs1，先求出第一种情况的sum以及每个颜色对应的贡献color[v[i]]
2. 遍历到rt的每个子树，我们先对每个子树清除答案，目的就是清除那些不经过根节点的路径的贡献
3. dfs2，求出第二种情况的贡献
4. 恢复该子树的答案
5. 遍历rt的每个子树，重复上面的4个操作
6. logn层的点分治即可

这就是全部了具体操作看代码比较好

代码实现

#include
using namespace std;
#define ll long long
const int N = 1e5+10;
int head[N],tot,n;
int rt,sum,maxp[N],siz[N],cnt[N],V[N],color[N],much,maxx;
ll ans[N],num;
bool vis[N];
struct edge{
    int v,nex,w;
    edge(){}
    edge(int v, int nex, int w) : v(v), nex(nex), w(w) {}
}edges[N<<1];
void add(int u,int v,int w){
    edges[tot]=edge(v,head[u],w);
    head[u]=tot++;
}
void dfs1(int u,int fa){
    siz[u]=1;
    cnt[V[u]]++;
    for(int i = head[u];~i;i=edges[i].nex){
        int v = edges[i].v;
        if(vis[v]||v==fa)continue;
        dfs1(v,u);
        siz[u]+=siz[v];
    }
    //首次出现
    if(cnt[V[u]]==1){
        sum+=siz[u];
        color[V[u]]+=siz[u];
    }
    cnt[V[u]]--;
}
//找重心
void getrt(int u,int fa){
    siz[u]=1,maxp[u]=0;
    for(int i=head[u];~i;i=edges[i].nex){
        int v =edges[i].v;
        if(vis[v]||v==fa)continue;
        getrt(v,u);
        siz[u]+=siz[v];
        if(maxp[u]<siz[v])maxp[u]=siz[v];
    }
    maxp[u]=max(maxp[u],maxx-siz[u]);
    if(maxp[u]<maxp[rt])rt=u;
}
void change(int u,int fa,int value){
    cnt[V[u]]++;
    for(int i=head[u];~i;i=edges[i].nex){
        int v = edges[i].v;
        if(v==fa||vis[v])continue;
        change(v,u,value);
    }
    if(cnt[V[u]]==1){
        sum += 1ll*siz[u]*value;
        color[V[u]]+=1ll*siz[u]*value;
    }
    cnt[V[u]]--;
}
void dfs2(int u,int fa){
    cnt[V[u]]++;
//    从u~p上的颜色数统计
    if(cnt[V[u]]==1){
        sum-=color[V[u]];
        num++;
    }
    ans[u]+=sum+num*much;
    for(int i=head[u];~i;i=edges[i].nex){
        int v = edges[i].v;
        if(vis[v]||v==fa)continue;
        dfs2(v,u);
    }
//    第一次出现的
    if(cnt[V[u]]==1){
        sum+=color[V[u]];
        num--;
    }
    cnt[V[u]]--;
}
void clear(int u,int fa){
    cnt[V[u]]=0;
    color[V[u]]=0;
    for(int i=head[u];~i;i=edges[i].nex){
        int v = edges[i].v;
        if(vis[v]||v==fa)continue;
        clear(v,u);
    }
}
void solve(int u,int fa){
    dfs1(u,fa);
    ans[u]+=sum-color[V[u]]+siz[u];
    for(int i=head[u];~i;i=edges[i].nex){
        int v = edges[i].v;
        if(vis[v]||v==fa)continue;
//        清除子树的贡献，因为计算的是经过根节点的路径
//        而我们计算当前子节点的时候需要清除点分治的第二种情况
        cnt[V[u]]++;
        sum-=siz[v];
        color[V[u]]-=siz[v];
        change(v,u,-1);
        cnt[V[u]]--;
        much=siz[u]-siz[v];
//        计算子树的贡献
        dfs2(v,u);
//        计算完了之后再加上当前子树的贡献
        cnt[V[u]]++;
        sum+=siz[v];
        color[V[u]]+=siz[v];
        change(v,u,1);
        cnt[V[u]]--;
    }
    sum = 0,num = 0;
//    删除当前点
    clear(u,fa);
}
//点分治
void divide(int u,int fa){
    vis[u]=1;
    solve(u,fa);
    for(int i=head[u];~i;i=edges[i].nex){
        int v = edges[i].v;
        if(vis[v])continue;
        maxp[rt=0]=maxx=siz[v];
        getrt(v,0);
        getrt(rt,0);
        divide(rt,0);
    }
}
int main(){
    memset(head,-1,sizeof(head));
    scanf("%d",&n);
    for(int i=1;i<=n;i++)scanf("%d",&V[i]);
    for(int i=0;i<n-1;i++){
        int u,v,w;
        scanf("%d %d",&u,&v);
        add(u,v,0);
        add(v,u,0);
    }
    maxp[0]=maxx=n;
    getrt(1,0);
    getrt(rt,0);
    divide(rt,0);
    for(int i=1;i<=n;i++)printf("%lld\n",ans[i]);
}