【jzoj5055】【GDOI2017模拟二试4.12】【树上路径】【点分治】

题目大意

给定一颗n个结点的无根树，树上的每个点有一个非负整数点权，定义一条路径的价值为路径上的点权和-路径的点权最大值。

给定参数p，我们想知道，有多少不同的树上简单路径，满足它的价值恰好是p的倍数。

注意：单点算作一个路径；u ≠ v时，（u,v）和（v,u）只算一次。

解题思路

直接点分即可。

code

#include<set>
#include<cmath>
#include<cstdio>
#include<cstring>
#include<algorithm>
#define LD double
#define LL long long
#define ULL unsigned long long
#define min(a,b) ((a<b)?a:b)
#define max(a,b) ((a>b)?a:b)
#define fo(i,j,k) for(int i=j;i<=k;i++)
#define fd(i,j,k) for(int i=j;i>=k;i--)
#define fr(i,j) for(int i=begin[j];i;i=next[i])
using namespace std;
int const mn=1e5+9,mm=2*1e5+9,inf=1e9;
int n,p,gra,top,a[mn],vis[mn],begin[mn],size[mn],to[mm],next[mm],
    cnt[10000000+9];LL ans;
void insert(int u,int v){
    to[++gra]=v;
    next[gra]=begin[u];
    begin[u]=gra;
}
void gsize(int now,int pre){
    size[now]=1;
    fr(i,now)if((!vis[to[i]])&&to[i]!=pre){
        gsize(to[i],now);
        size[now]+=size[to[i]];
    }
}
int groot(int now,int pre,int nsize){
    fr(i,now)if((!vis[to[i]])&&to[i]!=pre){
        int tmp=groot(to[i],now,nsize);
        if(tmp)return tmp;
    }
    if(size[now]>nsize/2)return now;
    return 0;
}
struct rec{
    LL sum;int max;
};
rec b[mn];
void dfs(int now,int pre,LL nsum,int nmax){
    b[++top].sum=nsum;
    b[top].max=nmax;
    fr(i,now)if((!vis[to[i]])&&to[i]!=pre)
        dfs(to[i],now,nsum+a[to[i]],max(nmax,a[to[i]]));
}
bool cmp(rec x,rec y){
    return x.max<y.max;
}
void calc(int now,LL nsum,int nmax,int tag,int tmp){
    top=0;
    dfs(now,0,nsum,nmax);
    sort(b+1,b+top+1,cmp);
    fo(i,1,top){
        ans+=tag*cnt[(b[i].sum-b[i].max)%p];
        cnt[((p-(b[i].sum-tmp))%p+p)%p]++;
    }
    fo(i,1,top)cnt[((p-(b[i].sum-tmp))%p+p)%p]--;
}
void solve(int now){
    gsize(now,0);
    now=groot(now,0,size[now]);
    calc(now,a[now],a[now],1,a[now]);vis[now]=1;
    fr(i,now)if(!vis[to[i]]){
        calc(to[i],a[now]+a[to[i]],max(a[now],a[to[i]]),-1,a[now]);
        solve(to[i]);
    }
}
int main(){
    //freopen("path.in","r",stdin);
    //freopen("path.out","w",stdout);
    freopen("d.in","r",stdin);
    freopen("d.out","w",stdout);
    scanf("%d%d",&n,&p);int u,v;
    fo(i,1,n-1){
        scanf("%d%d",&u,&v);
        insert(u,v);insert(v,u);
    }
    fo(i,1,n)scanf("%d",&a[i]);
    solve(1);
    printf("%lld\n",ans+n);
    return 0;
}