【JZOJ5055】【GDOI2017模拟二试4.12】树上路径

Description

给定一颗n个结点的无根树，树上的每个点有一个非负整数点权，定义一条路径的价值为路径上的点权和-路径的点权最大值。
给定参数p，我们想知道，有多少不同的树上简单路径，满足它的价值恰好是p的倍数。
注意：单点算作一个路径；u ≠ v时，（u,v）和（v,u）只算一次。

Data Constraint

对所有测试点，我们有：
n≤10^5,p≤10^7,val_i≤10^9

Solution

这是道树分治的题。我们找出重心的位置，每次从重心往四周遍历，找出每条到重心的路径的点权和%p和路径的点权最大值,然后将路径按点权最大值从小大大排序，用个桶维护当前的路径的点权和，每次在桶中查找路径的点权和-路径的点权最大值的数量。由于可能会算重，所以要先重心的每颗子树自己先搞一下，减去重复。

Code

#include<iostream>#include<cmath>#include<cstring>#include<cstdio>#include<algorithm>using namespace std;const int maxn=2e5+5,maxn1=1e7+5;struct code{    int mx,sum;}b[maxn];int first[maxn],last[maxn],next[maxn],a[maxn],size[maxn],mx[maxn];int n,i,t,j,k,l,m,x,y,z,num,p,ans,ln,s,cnt[maxn1][2],bz[maxn],fa[maxn];void lian(int x,int y){    last[++num]=y;next[num]=first[x];first[x]=num;}bool cmp(code x,code y){    return x.mx<y.mx;}void dg1(int x,int y){    int t,p=num;size[x]=1;mx[x]=0;    for (t=first[x];t;t=next[t]){        if (last[t]==y || bz[last[t]])continue;        b[++num].sum=(b[p].sum+a[last[t]])%m;        b[num].mx=max(a[last[t]],b[p].mx);        dg1(last[t],x);size[x]+=size[last[t]];mx[x]=max(mx[x],size[last[t]]);    }}int find(int x,int y){    int t,k;mx[x]=max(mx[x],p-size[x]);    if (mx[x]*2<=p||p==1) return x;    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 dg(int x){    int t,k;    bz[x]=1;num=0;    for (t=first[x];t;t=next[t]){        if (bz[last[t]]) continue;        k=num+1;        b[++num].mx=max(a[x],a[last[t]]);b[num].sum=(a[x]+a[last[t]])%m;        dg1(last[t],0);        sort(b+k,b+num+1,cmp);        for (i=k;i<=num;i++){            k=((b[i].sum-b[i].mx)%m+m)%m;            if (k) k=m-k;            if (cnt[k][0]==last[t]) ans-=cnt[k][1];            l=((b[i].sum-a[x])%m+m)%m;            if (cnt[l][0]!=last[t]) cnt[l][0]=last[t],cnt[l][1]=0;            cnt[l][1]++;        }    }    sort(b+1,b+num+1,cmp);    for (i=1;i<=num;i++){        k=((b[i].sum-b[i].mx)%m+m)%m;        if (k) k=m-k;else ans++;        if (cnt[k][0]==x) ans+=cnt[k][1];        l=((b[i].sum-a[x])%m+m)%m;        if (cnt[l][0]!=x) cnt[l][0]=x,cnt[l][1]=0;        cnt[l][1]++;    }    for (t=first[x];t;t=next[t]){        if (bz[last[t]])continue;p=size[last[t]];        k=find(last[t],x);        dg(k);    }}int main(){    freopen("path.in","r",stdin);freopen("path.out","w",stdout);    scanf("%d%d",&n,&m);    for (i=1;i<n;i++)        scanf("%d%d",&x,&y),lian(x,y),lian(y,x);    for (i=1;i<=n;i++)        scanf("%d",&a[i]);num=0;    dg1(1,0);p=n;    k=find(1,0);    dg(k);    ans+=n;    printf("%d\n",ans);}