【JZOJ3872】【NOIP2014八校联考第4场第1试10.19】圣诞树(tree)

Description

圣诞节到了，小可可送给小薰一棵圣诞树。这棵圣诞树很奇怪，它是一棵多叉树，有n个点，n-1条边。它的每个结点都有一个权值。小可可和小薰想用这棵树玩一个游戏。
定义(s,e)为树上从s到e的简单路径，我们可以记下在这条路径上经过的结点，定义这个结点序列为S(s,e)。
我们按照如下方法定义这个序列S(s,e)的权值G(S(s,e))：假设这个序列中结点的权值为Z0，Z1，…，Z(L-1)，其中L为序列的长度，我们定义G(S(s,e))=Z0 × k^0 + Z1 × k^1 + … + Z(L-1) × k^(L-1)。
如果路径(s,e)满足G(S(s,e)) ≡ x (mod y) ，那么这条路径属于小可可，否则这条路径属于小薰。小可可和小薰很显然不希望这个游戏变得那么简单。小薰认为如果路径(p1,p2)和(p2,p3)都属于他，那么路径(p1,p3)也属于他，反之如果路径(p1,p2)和(p2,p3)都属于小可可，那么路径(p1,p3)也属于小可可。然而这个性质并不总是正确的。所以小薰想知道到底有多少三元组(p1,p2,p3)满足这个性质。
小薰表示她看一眼就知道这道题怎么做了。你会吗？

Data Constraint

对于20%的数据，n ≤ 200；
对于50%的数据，n ≤ 10^4；
对于100%的数据，1 ≤ n ≤ 10^5，2 ≤ y ≤ 10^9，1 ≤ k ≤ y，0 ≤ x < y。

Solution

我们设出in0[i]表示i出发的最后能做到≡ x (mod y)的个数，in1[i]表示i出发的最后不能做到≡ x (mod y)的个数。out0[i],out1[i]同理。直接求满足性质的数目很难求，所以我们正难则反，求不满足性质的个数。那么就有$t=\sum_{i=1}^{n}{in1[i]*in0[i]*2+out0[i]*out1[i]*2+in0[i]*out1[i]+in1[i]*out0[i]}$
由于每个三角形最后被算了两遍，所以$ans=n^3-t/2$。求in1[i],out1[i]很难求，但我们发现in1[i]+in0[i]=out0[i]+out1[i]=n，所以现在只要求出in0[i],out0[i]即可。
我们可以用树分治来解决。每次跳到当前子树的重心，然后把这棵子树暴力遍历一遍，求出子树内每个点x到重心y的G(S(x,y))和G(S(y,x))。那么现在考虑一下合并。显然我们想要求的还有经过y的两两子节点的G值。那么就有$G(S(x,x1))=G(S(x,y))+G(S(y,x1))*k^{lenS(x,y)}$。我们要满足G(S(x,x1))≡ x (mod y)，那么$G(S(y,x1))={X-G(S(x,y))\over{k^{lenS(x,y)}}}$。我们只要在事先求出的G(S(y,x))中二分查找一下即可求出对于该G(S(x,y))满足G(S(x,x’))≡ x (mod y)的数量，加在out0[x]上，同时在对应的y上加1。问题来了，怎样在在对应的y上加1用是最短呢？假定[a,b]的区间均要加1，我们直接在a加1，b+1减1，最后扫一遍即可。同时，由于两个点x,y可能来自同一棵子树，所以我们要在重心的每个儿子各自再遍历一遍，做一遍上述操作，减去这些情况。最后再跳到重心的儿子上，找出那个子树的的重心……时间复杂度$O(Nlog^2N)$

Code

#include<iostream>
#include<cmath>
#include<cstring>
#include<cstdio>
#include<algorithm>
#define ll long long
using namespace std;
const ll maxn=1e5+5;
ll first[maxn],last[2*maxn],next[2*maxn],a[maxn],f[maxn],g[maxn],size[maxn],bz[maxn];
ll mx[maxn],h[maxn],in[maxn],out[maxn],d[maxn],er[maxn],f1[maxn];
struct code{
    ll a,b,c;
}b[maxn],c[maxn];
ll n,m,i,t,j,k,l,y,x,z,num,p,q,num1,ln,ans;
void lian(ll x,ll y){
    last[++num]=y;next[num]=first[x];first[x]=num;
}
bool cmp(code x,code y){
    return x.a<y.a;
}
ll mi(ll x,ll y){
    if (y==1) return x;
    ll t=mi(x,y/2);
    if (y%2) return t*t%p*x%p;return t*t%p;
}
void dg(int x,int y){
    int t;size[x]=1;mx[x]=0;
    for (t=first[x];t;t=next[t]){
        if (last[t]==y || bz[last[t]]) continue;
        dg(last[t],x);
        mx[x]=max(mx[x],size[last[t]]);
        size[x]+=size[last[t]];
    }
}
int dg1(int x,int y,int z){
    int t,k;mx[x]=max(mx[x],size[z]-size[x]);
    if (mx[x]<=size[z]/2) return x;
    for (t=first[x];t;t=next[t]){
        if (last[t]==y || bz[last[t]]) continue;
        k=dg1(last[t],x,z);
        if (k) return k;
    }
    return 0;
}
void dg2(int x,int y,ll z){
    int t,k=num;
    for (t=first[x];t;t=next[t]){
        if (last[t]==y || bz[last[t]]) continue;
        b[++num].a=(b[k].a+z*a[last[t]])%p;b[num].b=b[k].b+1;
        c[num].a=(c[k].a*m+a[last[t]])%p;c[num].b=c[k].b+1;
        b[num].c=c[num].c=last[t];
        dg2(last[t],x,z*m%p);
    }
}
void out0(ll x){
    int i,j,t,k,l,r,mid;b[num+1].a=p+1;b[0].a=-1;
    for (i=1;i<=num;i++){
        l=0;
        r=num+1;
        y=(q-c[i].a+p)%p*f1[c[i].b]%p;
        while (l<r){
            mid=(l+r)/2;
            if (b[mid].a>=y) r=mid;
            else l=mid+1;
        }
        t=l;
        l=0;
        r=num+1;
        while (l<r){
            mid=(l+r+1)/2;
            if (b[mid].a>y) r=mid-1;
            else l=mid;
        }
        k=l;
        if (t<=k) d[k+1]-=x,d[t]+=x,out[c[i].c]+=x*(k-t+1);
    }
    for (i=1;i<=num;i++)
        d[i]+=d[i-1],in[b[i].c]+=d[i];
    for (i=1;i<=num+2;i++)
        d[i]=0;
}
void make(ll x){
    int t,k,i;
    dg(x,0);
    h[x]=dg1(x,0,x);bz[h[x]]=1;num=1,b[1].a=0,c[1].a=a[h[x]],b[1].b=0,c[1].b=1;b[1].c=c[1].c=h[x];
    dg2(h[x],0,1);
    sort(b+1,b+num+1,cmp);sort(c+1,c+num+1,cmp);
    out0(1);
    for (t=first[h[x]];t;t=next[t])
        if (!bz[last[t]]){
            num=1,b[1].a=0,c[1].a=a[last[t]],b[1].b=0,c[1].b=1;b[1].c=c[1].c=last[t];
            dg2(last[t],0,1);
            for (i=1;i<=num;i++)
                c[i].a=(c[i].a+a[h[x]]*er[c[i].b])%p,b[i].b++,b[i].a=(b[i].a*m+a[last[t]])%p,c[i].b++;
            sort(b+1,b+num+1,cmp);sort(c+1,c+num+1,cmp);
            out0(-1);   
        }
    for (t=first[h[x]];t;t=next[t])
        if (!bz[last[t]]) make(last[t]);    
}
int main(){
//  freopen("data.in","r",stdin);
    scanf("%lld%lld%lld%lld",&n,&p,&m,&q);er[0]=1;
    for (i=1;i<=n;i++)
        er[i]=er[i-1]*m%p,f1[i]=mi(er[i],p-2);
    for (i=1;i<=n;i++)
        scanf("%lld",&a[i]);
    for (i=1;i<n;i++)
        scanf("%lld%lld",&x,&y),lian(x,y),lian(y,x);
    make(1);memset(bz,0,sizeof(bz));t=0;ans=n*n*n;
    for (i=1;i<=n;i++){
        x=n-in[i];y=n-out[i];
        t+=2*(x*in[i]+y*out[i])+in[i]*y+x*out[i];
    }
    ans-=t/2;
    printf("%lld\n",ans);
}