P2664 树上游戏 点分治

　　

 

 

树上点对问题很容易就想到点分治   

点分治已经用掉一个log了  那么我们的问题就转化为：

在 o(n) 时间内求出 以该根为lca的所有的点对的答案！

1、对于树上的一个i  如果i的颜色在i到根的链上是第一次出现 那么可以对树上其他的点j （不考虑特殊情况）产生贡献siz[i] （j与子树i的每一个点都产生贡献1）

2、dfs求出sum（整棵树所有的贡献 和color[i] 这个颜色所有的贡献） 

3、显然  根的答案为  ans[root]+=sum 也就是一端在根上的链 

4、遍历子树 统计答案  

　 但是有一些不是树上最短路径的贡献不能被统计进去  所以要在该子树统计答案的时候  删除该子树的影响  然后再回溯回去

5、设Ｘ＝sigma color[j 到根上（不包括根）的所有颜色] 由于这些颜色已经出现过，我们不能在该子树外计算其贡献）

　　设num为ｊ到根上（不包括根）的颜色数

　　设Ｙ为size[root]-size[该子树]（即所有其他子树＋根的点数）

　　则ans[j]+＝sum-X+num*Y;

 

这样一道黑题就大功告成啦

#include<bits/stdc++.h>
using namespace std;
#define rep(i,a,b) for(int i=(a);i<=(b);i++)
#define repp(i,a,b) for(int i=(a);i>=(b);--i)
#define ll long long
#define see(x) (cerr<<(#x)<<'='<<(x)<<endl)
#define inf 0x3f3f3f3f
#define CLR(A,v)  memset(A,v,sizeof A)
//
const int N=2e6+10;
int head[N],pos,n,m,v[N],siz[N],maxson[N],root,color[N],cnt[N],siz_tree,vis[N],a,b;
ll sum,ans[N],much,num;
struct Edge{int to,nex;}edge[N<<1];
void add(int a,int b){edge[++pos]=(Edge){b,head[a]};head[a]=pos;}
void getroot(int x,int fa)
{
    siz[x]=1;maxson[x]=0;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v]||v==fa)continue;
        getroot(v,x);siz[x]+=siz[v];
        maxson[x]=max(maxson[x],siz[v]);
    }
    maxson[x]=max(maxson[x],siz_tree-siz[x]);
    if(maxson[x]<maxson[root])root=x;
}
void dfs1(int x,int fa)
{
    cnt[v[x]]++;siz[x]=1;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(v==fa||vis[v])continue;
        dfs1(v,x);siz[x]+=siz[v];
    }
    if(cnt[v[x]]==1)
    {
        sum+=siz[x];
        color[v[x]]+=siz[x];
    }
    cnt[v[x]]--;
}
void dfs2(int x,int fa)
{
    cnt[v[x]]++;
    if(cnt[v[x]]==1)sum-=color[v[x]],num++;
    ans[x]+=sum+num*much;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v]||v==fa)continue;
        dfs2(v,x);
    }
    if(cnt[v[x]]==1)sum+=color[v[x]],num--;
    cnt[v[x]]--;
}
void change(int x,int fa,int value)
{
    cnt[v[x]]++;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(v==fa||vis[v])continue;
        change(v,x,value);
    }
    if(cnt[v[x]]==1)
    {
        sum+=1ll*value*siz[x];
        color[v[x]]+=1ll*value*siz[x];
    }
    cnt[v[x]]--;
}
void clear(int x,int fa)
{
    cnt[v[x]]=color[v[x]]=0;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(v==fa||vis[v])continue;
        clear(v,x);
    }
}
void calc(int x)
{
    dfs1(x,0);
    ans[x]+=sum;
    for(int i=head[x];i;i=edge[i].nex)
    {
        int to=edge[i].to;
        if(vis[to])continue;
        cnt[v[x]]++;
        sum-=siz[to];
        color[v[x]]-=siz[to];
        change(to,x,-1);
        much=siz[x]-siz[to];
        cnt[v[x]]--;
        dfs2(to,x);

        cnt[v[x]]++;
        sum+=siz[to];
        color[v[x]]+=siz[to];
        change(to,x,1);//回溯
        cnt[v[x]]--;
    }
    sum=num=0;
    clear(x,0);
}
void solve(int x)
{
    if(vis[x])return ;vis[x]=1;
    calc(x);
    for(int i=head[x];i;i=edge[i].nex)
    {
        int v=edge[i].to;
        if(vis[v])continue;
        root=0;
        siz_tree=siz[v];
        getroot(v,x);
        solve(root);
    }
}
int main()
{
    scanf("%d",&n);
    rep(i,1,n)scanf("%d",&v[i]);
    rep(i,1,n-1)scanf("%d%d",&a,&b),add(a,b),add(b,a);
    siz_tree=maxson[0]=n;
    root=0;
    getroot(1,0);
    solve(root);
    rep(i,1,n)printf("%lld\n",ans[i]);
    return 0;
}

View Code

 

转载于:https://www.cnblogs.com/bxd123/p/11552611.html