[JZOJ5055]树上路径

题目大意

给定一棵$n$个节点的无根树，每个点都有一个非负整数的权值${val}_i$，定义一条路径的价值为路径上的点权和减去路径的点权最大值。
给定参数$p$，请求出树上有多少条价值是$p$的倍数的路径。
注意：单点也算路径。并且路径$(u,v)$和$(v,u)$只算一次。

$1\le n\le10^5,1\le p\le10^7,0\le {val}_i\le10^9$

---

题目分析

很直观的想法是点剖。
考虑对于一个分治重心，我依次处理其各个子树。对于一条重心到一个点的路径，假设其权值和是${sum}_2$，最大值是${mx}_2$，在以前遍历过的子树中，如果有一条权值和是${sum}_1$，最大值是${mx}_1$的路径满足条件，当且仅当一下两个条件其一成立：
$\bullet\ {mx}_1\le{mx}_2$，且${sum}_1+{sum}_2\equiv{mx}_2\pmod p$
$\bullet\ {mx}_1>{mx}_2$，且${sum}_1+{sum}_2\equiv{mx}_1\pmod p$
把后面的式子瞎jb化一下，变成只和$i$有关的形式，就可以开很多棵一棵动态开点的权值线段树来查询满足前面条件的点数了。
时间复杂度$\mathrm O(n\log^2n)$，常数有点大。
还存在某排序加删掉重复部分的算法，时间复杂度相同，但是常数小很多。

---

代码实现

#include <iostream>
#include <cstdio>
#include <cctype>

using namespace std;

typedef long long LL;

int read()
{
    int x=0,f=1;
    char ch=getchar();
    while (!isdigit(ch)) f=ch=='-'?-1:f,ch=getchar();
    while (isdigit(ch)) x=x*10+ch-'0',ch=getchar();
    return x*f;
}

const int N=100050;
const int E=N<<1;
const int P=10000000;
const int V=1000000000;
const int LGV=31;
const int S=N*LGV;

struct segment_tree
{
    int son[S][2];
    int size[S];
    int tot;

    void init(){tot=0;}

    int newnode()
    {
        size[++tot]=0,son[tot][0]=son[tot][1]=0;
        return tot;
    }

    void insert(int &rt,int x,int l,int r)
    {
        if (!rt) rt=newnode();
        ++size[rt];
        if (l==r) return;
        int mid=l+r>>1;
        if (x<=mid) insert(son[rt][0],x,l,mid);
        else insert(son[rt][1],x,mid+1,r);
    }

    int query(int rt,int st,int en,int l,int r)
    {
        if (!rt) return 0;
        if (st==l&&en==r) return size[rt];
        int mid=l+r>>1;
        if (en<=mid) return query(son[rt][0],st,en,l,mid);
        else if (mid+1<=st) return query(son[rt][1],st,en,mid+1,r);
        else return query(son[rt][0],st,mid,l,mid)+query(son[rt][1],mid+1,en,mid+1,r);
    }
}t[2];

int last[N],fa[N],que[N],size[N],val[N];
int tmp[N][2],temp[N][2];
int tov[E],nxt[E];
int root[P][2];
bool vis[N];
int n,p,head,tail,tot,tmps,temps;
LL ans;

void insert(int x,int y){tov[++tot]=y,nxt[tot]=last[x],last[x]=tot;}

int core(int og)
{
    int i,x,y,ret,rets=n,tmp;
    for (head=0,fa[que[tail=1]=og]=0;head<tail;)
        for (size[x=que[++head]]=1,i=last[x];i;i=nxt[i])
            if ((y=tov[i])!=fa[x]&&!vis[y]) fa[que[++tail]=y]=x;
    for (head=tail;head>1;--head) size[fa[que[head]]]+=size[que[head]];
    for (head=1;head<=tail;++head)
    {
        for (i=last[x=que[head]],tmp=size[og]-size[x];i;i=nxt[i])
            if ((y=tov[i])!=fa[x]&&!vis[y]) tmp=max(tmp,size[y]);
        if (rets>tmp) rets=tmp,ret=x;
    }
    return ret;
}

void dfs(int x,int sum,int mx)
{
    (sum+=val[x])%=p,mx=max(mx,val[x]);
    ans+=t[0].query(root[(mx%p-sum+p)%p][0],0,mx,0,V);
    ans+=t[1].size[root[sum][1]]-t[1].query(root[sum][1],0,mx,0,V);
    tmp[++tmps][0]=sum,tmp[tmps][1]=mx;
    for (int i=last[x],y;i;i=nxt[i])
        if ((y=tov[i])!=fa[x]&&!vis[y]) fa[y]=x,dfs(y,sum,mx);
}

void solve(int x)
{
    int c=core(x);
    fa[c]=0,t[0].init(),t[1].init();
    for (int i=last[c];i;i=nxt[i])
        if (!vis[x=tov[i]])
        {
            fa[x]=c,dfs(x,0,0);
            for (int sum,mx;tmps;--tmps)
            {
                sum=(tmp[tmps][0]+val[c]%p)%p,mx=max(tmp[tmps][1],val[c]),ans+=sum==mx%p;
                t[0].insert(root[sum][0],mx,0,V);
                t[1].insert(root[(mx%p-sum+p)%p][1],mx,0,V);
                temp[++temps][0]=sum,temp[temps][1]=mx;
            }
        }
    for (int sum,mx;temps;--temps)
    {
        sum=temp[temps][0],mx=temp[temps][1];
        root[sum][0]=root[(mx%p-sum+p)%p][1]=0;
    }
    vis[c]=1;
    for (int i=last[c],y;i;i=nxt[i])
        if (!vis[y=tov[i]]) solve(y);
}

int main()
{
    freopen("path.in","r",stdin),freopen("path.out","w",stdout);
    n=read(),p=read();
    for (int i=1,x,y;i<n;++i) x=read(),y=read(),insert(x,y),insert(y,x);
    for (int i=1;i<=n;++i) val[i]=read();
    solve(1);
    printf("%lld\n",ans+n);
    fclose(stdin),fclose(stdout);
    return 0;
}