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;
}