【GDSOI2017模拟4.13】采蘑菇（点剖||线段树）

Description

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

Solution

这道题目可以直接用线段树处理换根的问题，还可以直接用点剖来做。
首先我们要处理出这个子树对这个root节点的答案，我们设an[i]表示这个子树对i的贡献（所有它子节点到i路径颜色数总和）
然后我们根据树分治的套路，处理出所有经过root节点到x的所有路径的贡献。
假设从root到x路径上的颜色数为b，那么要从root的答案转移到x首先就要
an[root]−an[x]+(size[root]−size[x])∗b （除了x的子树节点都转移b的贡献）
但是有可能别的节点到root之前已经有b颜色集中的某一部分了，所以要把这些点给排掉，设aw[i]为有多少个x节点到root的路径上包含i颜色的， c=∑i∈baw[i]
那么还要减去(c-b*(size[x]))，因为size[x]不属于当前考虑的去重部分，且他们一定包含b个颜色。
所以每个
ans[x]+=an[root]−an[x]+(size[root]−size[x])∗b−(c−b∗(size[x]))
=an[root]−an[x]+size[root]∗b−c
根据树分治的套路但是还可能有root的直系儿子并未x的祖先y，y这个子树也会又贡献，那么就会有问题，所以设上面对所有x的操作的递归为dfs(x,1)，1为正贡献，然后对root的所有直系儿子做一次dfs(y,-1)，-1为负贡献，这里dfs时的桶里要包含a[x]这个颜色，aw也要，注意这两个数组要重新求。

Code

#include
#include
#include
#include
#include
#define fo(i,a,b) for(i=a;i<=b;i++)
#define fod(i,a,b) for(i=a;i>=b;i--)
#define rep(i,a) for(i=first[a];i;i=next[i])
const int maxn=3e5+7;
typedef long long ll;
int first[maxn*2],last[maxn*2],next[maxn*2],num,a[maxn];
int i,j,k,l,n,m,size[maxn],z,dian,fa[maxn],aw[maxn][2],t[maxn],dan[maxn],c,o,b,zf,tim;
int sum1,sum2;
bool bz[maxn];
ll ans[maxn],an[maxn];
bool az;
void add(int x,int y){
    last[++num]=y,next[num]=first[x],first[x]=num;
}
void dfs1(int x,int y){
    int i;
    size[x]=1;dan[x]=dan[y];
    if(!t[a[x]])dan[x]++;t[a[x]]++;if(!az)an[x]=dan[x];
    rep(i,x)if(!bz[last[i]]&&last[i]!=y){
        dfs1(last[i],x),size[x]+=size[last[i]];
        if(!az)an[x]+=an[last[i]];
    }
    if(aw[a[x]][1]!=tim)aw[a[x]][1]=tim,aw[a[x]][0]=0;
    t[a[x]]--;
    if(!t[a[x]]&&z)aw[a[x]][0]+=size[x];
}
void dfs2(int x,int y){
    int i;bool az=1;
    rep(i,x)if(!bz[last[i]]&&last[i]!=y){
        if(size[last[i]]>dian/2)az=0;
        dfs2(last[i],x);
    }
    if(az&&dian-size[x]<=dian/2)z=x;
}
void dfs3(int x,int y){
    int i;
    if(aw[a[x]][1]!=tim)aw[a[x]][1]=tim,aw[a[x]][0]=0;
    if(x!=z&&!t[a[x]])c+=aw[a[x]][0],b++;t[a[x]]++;
    if(x!=z)ans[x]+=(size[z]*b+an[z]-an[x]-c)*zf;
    rep(i,x){
        if(!bz[last[i]]&&last[i]!=y)dfs3(last[i],x);
    }
    if(aw[a[x]][1]!=tim)aw[a[x]][1]=tim,aw[a[x]][0]=0;
    t[a[x]]--;if(x!=z&&!t[a[x]])c-=aw[a[x]][0],b--;
}
void fen(int x){
    int i;az=0;
    z=0;dfs1(x,0);dian=size[x];dfs2(x,0);x=z;
    bz[x]=1;++tim;dfs1(x,0);
    t[a[x]]++;ans[x]+=an[x];
    rep(i,x){
        if(bz[last[i]])continue;
        sum1=sum2=0;zf=1;z=x;b=1,c=size[x];
        dfs3(last[i],x);
    }
    rep(i,x){
        if(bz[last[i]])continue;
        sum1=sum2=0;zf=1;z=x;b=1,c=size[x];
        ++tim;
        if(aw[a[x]][1]!=tim)aw[a[x]][1]=tim,aw[a[x]][0]=0;
        aw[a[x]][0]=size[last[i]];z=1;az=1;
        dfs1(last[i],0);
        z=last[i],b=1,zf=-1,c=size[last[i]];
        dfs3(last[i],x);
    }
    t[a[x]]--;
    rep(i,x){
        if(!bz[last[i]])fen(last[i]);
    }
}
int main(){
    freopen("mushroom.in","r",stdin);
    freopen("mushroom.out","w",stdout);
    scanf("%d",&n);
    fo(i,1,n)scanf("%d",&a[i]);
    fo(i,1,n-1){
        scanf("%d%d",&k,&l);
        add(k,l),add(l,k);
    }
    fen(n/2);
    fo(i,1,n)printf("%d\n",ans[i]);
}