【JZOJ5058】【GDSOI2017模拟4.13】采蘑菇

Description

A君住在魔法森林里，魔法森林可以看做一棵n个结点的树，结点从1~n编号。树中的每个结点上都生长着蘑菇。蘑菇有许多不同的种类，但同一个结点上的蘑菇都是同一种类，更具体地，i号结点上生长着种类为c[i]的蘑菇。
现在A君打算出去采蘑菇，但他并不知道哪里的蘑菇更好，因此他选定起点s后会等概率随机选择树中的某个结点t作为终点，之后从s沿着(s,t)间的最短路径走到t.并且A君会采摘途中所经过的所有结点上的蘑菇。
现在A君想知道，对于每一个结点u，假如他从这个结点出发，他最后能采摘到的蘑菇种类数的期望是多少。为了方便，你告诉A君答案*n的值即可。

Data Constraint

30%的数据：n <= 2000
另有20%的数据：给出的第i条边为{i,i+1}
另有20%的数据：蘑菇的种类最多3种
100%的数据：1 <= n <= 3*10^5 , 0 <= c[i] <= n

Solution

这又是一道树分治的题。但我用线段树加换根打过了。我们维护颜色棵线段树维护区间[l,r]中每个点到当前根的路径上有这个颜色的点的数量。再换根时，我们先撤销根的父亲原来对根的子树的影响，在加上当前根对全图的影响，统计答案变化值+父亲作为根即可，时间复杂度O(NlogN)。

Code

#include<iostream>
#include<cmath>
#include<cstring>
#include<cstdio>
#include<algorithm>
using namespace std;
const int maxn=6e5+5,maxn1=3e5+5;
struct code{
    int l,r,bz,num;
}f[maxn*25];
int first[maxn],last[maxn],next[maxn],a[maxn1],d[maxn1],b[maxn1],v[maxn1];
int n,i,t,j,l,x,y,z,num,num1,dfn[maxn1],size[maxn1],fa[maxn1];
long long k,ans[maxn1];
void lian(int x,int y){
    last[++num]=y;next[num]=first[x];first[x]=num;
}
void insert(int &v,int l,int r,int x,int y,int z){
    int mid=(l+r)/2;
    if (x>y) return;
    if (!v) v=++num1,f[v].num=0;
    if (l!=r){
        if (!f[v].l) f[v].l=++num1;
        if (!f[v].r) f[v].r=++num1;
    }
    if (l>=x && r<=y){
        if (l==1 && r==n) k=-f[v].num;
        f[v].bz+=z;
        if (f[v].bz>0) f[v].num=r-l+1;
        else if (l!=r)f[v].num=f[f[v].l].num+f[f[v].r].num;
        else f[v].num=0;
        if (l==1 && r==n) k+=f[v].num;
        return;
    }
    if (l<=y && mid>=x) insert(f[v].l,l,mid,x,y,z);
    if (mid<y && r>=x) insert(f[v].r,mid+1,r,x,y,z);
    if (l==1 && r==n) k-=f[v].num;
    if (f[v].bz>0) f[v].num=r-l+1;
    else f[v].num=f[f[v].l].num+f[f[v].r].num;
    if (l==1 && r==n) k+=f[v].num;
}
void bfs(){
    int i=0,j=1,x,t;v[1]=1;fa[1]=1;
    while(i<j){
        x=v[++i];size[x]=1;
        for (t=first[x];t;t=next[t]){
            if (fa[last[t]])continue;
            v[++j]=last[t];fa[v[j]]=x;
        }
    }
    fa[1]=0;
    for (j=n;j>=1;j--)
        size[fa[v[j]]]+=size[v[j]];
    dfn[1]=1;
    for (i=1;i<=n;i++){
        x=v[i];
        k=1;
        for (t=first[x];t;t=next[t])
            if (last[t]!=fa[x])dfn[last[t]]=dfn[x]+k,k+=size[last[t]];
    }
    for (j=n;j>=1;j--)
        x=v[j],insert(d[a[x]],1,n,dfn[x],dfn[x]+size[x]-1,1);
    for (x=1;x<=n;x++)
        v[dfn[x]]=x;
}
int main(){
    freopen("mushroom.in","r",stdin);freopen("mushroom.out","w",stdout);
    scanf("%d",&n);
    for (i=1;i<=n;i++)
        scanf("%d",&a[i]);
    for (i=1;i<n;i++)
        scanf("%d%d",&x,&y),lian(x,y),lian(y,x);
    num=0;
    bfs();
    for (i=0;i<=n;i++)
        k=f[d[i]].num,ans[1]+=k;
    for (i=1;i<=n;i++){
        x=v[i];y=fa[x];b[++b[0]]=x;
        if (y){
            k=0;
            insert(d[a[y]],1,n,dfn[y],dfn[y]+size[y]-1,-1);
            insert(d[a[y]],1,n,1,dfn[x]-1,1);insert(d[a[y]],1,n,dfn[x]+size[x],n,1);
            insert(d[a[x]],1,n,1,dfn[x]-1,1);insert(d[a[x]],1,n,dfn[x]+size[x],n,1);
            ans[x]=ans[y]+k;
        }
        while (dfn[b[b[0]]]+size[b[b[0]]]-1==dfn[v[i]]){
            x=b[b[0]--];y=fa[x];
            if (y){
                insert(d[a[x]],1,n,1,dfn[x]-1,-1);insert(d[a[x]],1,n,dfn[x]+size[x],n,-1);
                insert(d[a[y]],1,n,1,dfn[x]-1,-1);insert(d[a[y]],1,n,dfn[x]+size[x],n,-1);
                insert(d[a[y]],1,n,dfn[y],dfn[y]+size[y]-1,1);
            }
        }
    }
    for (i=1;i<=n;i++)
        printf("%lld\n",ans[i]);
}

顺便附上我打的树分治

#include<iostream>
#include<cstring>
#include<cstdio>
#include<cmath>
#include<algorithm>
#define ll long long
using namespace std;
const int maxn=7e5+5;
ll aw[maxn],an[maxn],size[maxn],ans[maxn],mx[maxn];
int first[maxn],last[maxn],next[maxn],a[maxn];
int n,i,t,j,k,l,x,y,z,num,p,bz[maxn],bz1[maxn],cnt[maxn],q,sum;
void lian(int x,int y){
    last[++num]=y;next[num]=first[x];first[x]=num;
}
void dg1(int x,int y){int t;size[x]=1;mx[x]=0;aw[a[x]]=0;
    for (t=first[x];t;t=next[t])
        if (last[t]!=y && !bz[last[t]])dg1(last[t],x),size[x]+=size[last[t]],mx[x]=max(mx[x],size[last[t]]);
}
int find(int x,int y){
    if (max(p-size[x],mx[x])*2<=p) return x;int t;
    for (t=first[x];t;t=next[t]){
        if (last[t]==y || bz[last[t]]) continue;
        k=find(last[t],x);
        if (k) return k;
    }return 0;
}
void dg2(int x,int y,ll sum){int t;
    if (bz1[a[x]]!=p) bz1[a[x]]=p,cnt[a[x]]=0,sum++,aw[a[x]]+=size[x];
    cnt[a[x]]++;an[x]=sum;
    for (t=first[x];t;t=next[t]){
        if (last[t]==y || bz[last[t]]) continue;
        dg2(last[t],x,sum);
        an[x]+=an[last[t]];
    }
    cnt[a[x]]--;if (!cnt[a[x]]) bz1[a[x]]=-1;
}
void dg3(int x,int y,ll c){int t;
    if (bz1[a[x]]!=p) bz1[a[x]]=p,cnt[a[x]]=0,num++,c+=aw[a[x]];
    cnt[a[x]]++;ans[x]+=z*(num*(q-size[x])+(sum-an[x])-(c-size[x]*num));
    for (t=first[x];t;t=next[t]){
        if (last[t]==y || bz[last[t]]) continue;
        dg3(last[t],x,c);
    }
    cnt[a[x]]--;if (!cnt[a[x]]) bz1[a[x]]=-1,num--;
}
void dg(int x){
    int t;bz[x]=1;p=x;
    dg1(x,0);
    dg2(x,0,0);
    ans[x]+=an[x];cnt[a[x]]=1,num=1;z=1;
    for (t=first[x];t;t=next[t]){
        if (bz[last[t]]) continue;
        q=size[x];p=last[t];sum=an[x];bz1[a[x]]=p;
        dg3(last[t],x,aw[a[x]]);
    }
    z=-1;
    for (t=first[x];t;t=next[t]){
        if (bz[last[t]]) continue;p=last[t];
        dg1(last[t],0);bz1[a[x]]=p,cnt[a[x]]=1,aw[a[x]]=size[last[t]];
        dg2(last[t],0,1);
        q=size[last[t]];sum=an[last[t]];num=0;
        dg3(last[t],x,0);
    }
    bz1[a[x]]=-1;
    for (t=first[x];t;t=next[t])
        if (!bz[last[t]])p=size[last[t]],k=find(last[t],0),dg(k);
}
int main(){
    freopen("mushroom.in","r",stdin);freopen("mushroom.out","w",stdout);
    scanf("%d",&n);
    for (i=1;i<=n;i++)
        scanf("%d",&a[i]);
    for (i=1;i<n;i++)
        scanf("%d%d",&x,&y),lian(x,y),lian(y,x);
    dg(n/2);
    for (i=1;i<=n;i++)
        printf("%lld\n",ans[i]);
}