【清华集训模拟】树

Description

给出一棵n个点的树，每个点有点权a[i]
求把这棵树划分成若干条不相交的路径，使得每条路径上的点权和均非负的方案数。
n<=1e5

Solution

先考虑Dp怎么写，设Fx表示以x为根的子树已经覆盖完成了.
转移两条链，把这两条链上挂着的F值乘起来就是答案。
那么这个东西要如何用数据结构优化？
我们以权值为下标建立线段树，线段树里的每个点维护的是其到当前所做的根x路径上，所有挂着的F值的乘积。
先继承重儿子的线段树，对于每个轻儿子暴力在线段树上查询。
然后做完之后需要打上区间乘法标记。
注意我们可能没有逆元所以需要用前缀后缀积规避乘法。
复杂度可以看成启发式合并，O(n log^2 n)

Code

#include <cstdio>
#include <cstring>
#include <algorithm>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fd(i,a,b) for(int i=a;i>=b;i--)
#define rep(i,a) for(int i=lst[a];i;i=nxt[i])
using namespace std;

typedef long long ll;

int read() {
    char ch;int sig=1;
    for(ch=getchar();ch<'0'||ch>'9';ch=getchar()) if (ch=='-') sig=-1;
    int x=ch-'0';
    for(ch=getchar();ch>='0'&&ch<='9';ch=getchar()) x=x*10+ch-'0';
    return x*sig;
}

const int N=1e5+5,Mo=1e9+7,inf=0x7fffffff;

int t[N<<1],nxt[N<<1],lst[N],l;
void add(int x,int y) {
    t[++l]=y;nxt[l]=lst[x];lst[x]=l;
}

int n,x,y,a[N],dfn[N],size[N],dist[N],ord[N],id[N];
int F[N],G[N],pre[N],suf[N],c[N],w[N];
int root[N],tr[N<<7],ls[N<<7],rs[N<<7],lazy[N<<7],tot;

bool cmp(int x,int y) {return dist[x]<dist[y];}
bool cmp1(int x,int y) {return size[x]>size[y];}

void back(int v,int z) {
    tr[v]=(ll)tr[v]*z%Mo;
    lazy[v]=(ll)lazy[v]*z%Mo;
}

void down(int v) {
    if (lazy[v]!=1) {
        back(ls[v],lazy[v]);
        back(rs[v],lazy[v]);
        lazy[v]=1;
    }
}

void insert(int &v,int l,int r,int x,int y) {
    down(v);
    tr[++tot]=(tr[v]+y)%Mo;lazy[tot]=1;
    ls[tot]=ls[v];rs[tot]=rs[v];v=tot;
    if (l==r) return;
    int mid=l+r>>1;
    if (x<=mid) insert(ls[v],l,mid,x,y);
    else insert(rs[v],mid+1,r,x,y);
}

int find(int v,int l,int r,int x,int y) {
    if (!v||x>y) return 0;
    if (l==x&&r==y) return tr[v];
    int mid=l+r>>1;down(v);
    if (y<=mid) return find(ls[v],l,mid,x,y);
    else if (x>mid) return find(rs[v],mid+1,r,x,y);
    else return (find(ls[v],l,mid,x,mid)+find(rs[v],mid+1,r,mid+1,y))%Mo;
}

int merge(int x,int y) {
    if (!x||!y) return x+y;
    down(x);down(y);
    ls[x]=merge(ls[x],ls[y]);
    rs[x]=merge(rs[x],rs[y]);
    tr[x]=(tr[ls[x]]+tr[rs[x]])%Mo;
    return x;
}

void dfs(int x,int y) {
    dfn[x]=++tot;size[x]=1;w[tot]=x;
    rep(i,x) 
        if (t[i]!=y) {
            dist[t[i]]=dist[x]+a[t[i]];
            dfs(t[i],x);
            size[x]+=size[t[i]];
        }
}

int binary(int x) {
    int l=1,r=n+1;
    while (l<r) {
        int mid=l+r>>1;
        if (dist[ord[mid]]>=x) r=mid;
        else l=mid+1;
    }
    return l;
}

void Dp(int x,int y) {
    rep(i,x) if (t[i]!=y) Dp(t[i],x);

    c[0]=0;
    rep(i,x) if (t[i]!=y) c[++c[0]]=t[i];
    sort(c+1,c+c[0]+1,cmp1);
    pre[0]=1;fo(i,1,c[0]) pre[i]=(ll)pre[i-1]*F[c[i]]%Mo;
    suf[c[0]+1]=1;fd(i,c[0],1) suf[i]=(ll)suf[i+1]*F[c[i]]%Mo;

    G[x]=1;
    fo(i,1,c[0]) {
        int y=c[i];
        G[x]=(ll)G[x]*F[y]%Mo;
        (F[x]+=(ll)find(root[y],1,n,binary(dist[x]-a[x]),n)*pre[i-1]%Mo*suf[i+1]%Mo)%=Mo;
        if (i!=1) {
            fo(j,dfn[y],dfn[y]+size[y]-1)
                (F[x]+=(ll)find(root[y],1,n,id[w[j]],id[w[j]])*find(root[x],1,n,binary(2*dist[x]-a[x]-dist[w[j]]),n)%Mo*suf[i+1]%Mo)%=Mo;
        }
        back(root[x],F[y]);
        back(root[y],pre[i-1]);
        root[x]=merge(root[x],root[y]);
    }
    if (a[x]>=0) (F[x]+=G[x])%=Mo;
    insert(root[x],1,n,id[x],G[x]);
}

int main() {
    freopen("tree.in","r",stdin);
    freopen("tree.out","w",stdout);
    n=read();
    fo(i,1,n) a[i]=read();
    fo(i,1,n-1) {
        x=read();y=read();
        add(x,y);add(y,x);
    }
    dist[1]=a[1];
    dfs(1,0);
    dist[0]=inf;
    fo(i,1,n) ord[i]=i;
    sort(ord+1,ord+n+1,cmp);
    fo(i,1,n) id[ord[i]]=i;
    Dp(1,0);
    printf("%d\n",F[1]);
    return 0;
}