[JZOJ4331] 【清华集训模拟】树
题目
题目大意
给你一棵带点权的树，求将树变成一堆不相交的链，而且这些链的权值和非负的方案数。
正解
显然这道题是个\(DP\)。
首先求个前缀和\(sum\)。
为了后面讲述方便，我这样设：\(f_{i,j}\)表示以\(i\)为根的子树，其中某条链从\(x\)伸出到\(i\)的方案数，而且\(sum_x=j\)。
还有设\(g_i\)表示以\(i\)为根的，没有伸出去的链的方案数。
显然有这样的转移：
\[\prod g_i\to f_{x,sum_x}\\
f_{y,j}\prod_{i\neq y} g_i\to f_x,j\]
\[f_{x,j} \to g_x \ (j-sum_{fa_x}\geq 0)\\
f_{y,j}f_{z,k}\prod_{i\neq y,i\neq z}g_i\to g_x \ (j+k-2sum_x+a_x\geq 0)
\]
如果直接这样搞肯定会爆炸。所以考虑用线段树来维护\(f\)。
由于可能会出现\(g\)值为\(0\)的情况，所以不能直接用逆元来搞。
要维护个前缀积和后缀积。
首先要求出重儿子，把重儿子作为第一个儿子，然后线段树合并之前也启发式合并。
具体来说，我们钦定\(j<k\)。在合并的时候（设前面子树合并出来的线段树为\(A\)，这个线段树为\(B\)）当前的儿子作为\(k\)，遍历\(B\)的所有叶子节点，并在\(A\)中区间询问。这时候记得要乘上后缀积。将询问出来的东西加在\(g_x\)中。
然后两个合并在一起。记得在合并之前，整个\(A\)乘子树的\(g_x\)，整个\(B\)乘前缀积
。搞完这个再合并。
最后你就会愉快地发现，所有子树合并之后就是上面第二行式子。这样只需要把第一行的加进去。第四行的式子已经计算完了，只需要再加上第一个式子就可以了。
然而这不是题解的做法，作为一个小蒟蒻，表示看不懂题解……
代码
using namespace std;
#include <cstdio>
#include <cstring>
#include <algorithm>
#include <cassert>
#define N 100010
#define INF 1000000000
#define mo 1000000007
int n;
int a[N];
struct EDGE{
int to;
EDGE *las;
} e[N*2];
int ne;
EDGE *last[N];
struct Node *null;
struct Node{
Node *l,*r;
int sum,tag;
inline void get_tag(int c){sum=(long long)sum*c%mo;tag=(long long)tag*c%mo;}
inline void pushdown(){
if (tag!=1){
l->get_tag(tag);
r->get_tag(tag);
tag=1;
}
}
inline void update(){sum=l->sum+r->sum;sum>=mo?sum-=mo:0;}
} d[N*40];
int cnt;
Node *rt[N];
inline Node *newnode(){return &(d[++cnt]={null,null,0,1});}
void add(Node *&t,int l,int r,int x,int c){
if (t==null)
t=newnode();
if (l==r){
(t->sum+=c)%=mo;
return;
}
t->pushdown();
int mid=l+r>>1;
if (x<=mid)
add(t->l,l,mid,x,c);
else
add(t->r,mid+1,r,x,c);
t->update();
}
int query(Node *t,int l,int r,int st,int en){
if (t==null)
return 0;
if (st<=l && r<=en)
return t->sum;
t->pushdown();
int mid=l+r>>1,res=0;
if (st<=mid)
res+=query(t->l,l,mid,st,en);
if (mid<en)
res+=query(t->r,mid+1,r,st,en);
return res>=mo?res-mo:res;
}
Node *merge(Node *a,Node *b){
if (a==null)
return b;
if (b==null)
return a;
a->pushdown(),b->pushdown();
a->l=merge(a->l,b->l);
a->r=merge(a->r,b->r);
a->sum+=b->sum;
a->sum>=mo?a->sum-=mo:0;
return a;
}
int calc(Node *t,int l,int r,Node *rt,int bor){
if (t==null)
return 0;
if (l==r)
return (long long)query(rt,-INF,INF,bor-l,INF)*t->sum%mo;
t->pushdown();
int mid=l+r>>1,res=0;
res+=calc(t->l,l,mid,rt,bor);
res+=calc(t->r,mid+1,r,rt,bor);
return res>=mo?res-mo:res;
}
int fa[N],sum[N],siz[N],hs[N];
int son[N],ns,pre[N],suc[N];
int g[N];
void dp(int x){
sum[x]=sum[fa[x]]+a[x];
siz[x]=1;
for (EDGE *ei=last[x];ei;ei=ei->las)
if (ei->to!=fa[x]){
fa[ei->to]=x;
dp(ei->to);
siz[x]+=siz[ei->to];
if (siz[ei->to]>siz[hs[x]])
hs[x]=ei->to;
}
if (!hs[x]){
rt[x]=newnode();
add(rt[x],-INF,INF,sum[x],1);
g[x]=(a[x]>=0);
return;
}
son[ns=1]=hs[x];
pre[0]=1,pre[1]=g[hs[x]];
for (EDGE *ei=last[x];ei;ei=ei->las)
if (ei->to!=fa[x] && ei->to!=hs[x]){
son[++ns]=ei->to;
pre[ns]=(long long)pre[ns-1]*g[ei->to]%mo;
}
suc[ns+1]=1;
for (int i=ns;i>=1;--i)
suc[i]=(long long)suc[i+1]*g[son[i]]%mo;
rt[x]=rt[hs[x]];
for (int i=2;i<=ns;++i){
g[x]=(g[x]+(long long)calc(rt[son[i]],-INF,INF,rt[x],sum[fa[x]]*2+a[x])*suc[i+1])%mo;
rt[son[i]]->get_tag(pre[i-1]);
rt[x]->get_tag(g[son[i]]);
rt[x]=merge(rt[x],rt[son[i]]);
}
add(rt[x],-INF,INF,sum[x],pre[ns]);
g[x]+=query(rt[x],-INF,INF,sum[fa[x]],INF);
g[x]>=mo?g[x]-=mo:0;
}
int main(){
freopen("tree.in","r",stdin);
freopen("tree.out","w",stdout);
scanf("%d",&n);
for (int i=1;i<=n;++i)
scanf("%d",&a[i]);
for (int i=1;i<n;++i){
int u,v;
scanf("%d%d",&u,&v);
e[ne]={v,last[u]};
last[u]=e+ne++;
e[ne]={u,last[v]};
last[v]=e+ne++;
}
null=d;
*null={null,null,0,0};
dp(n>>1);
printf("%d\n",g[n>>1]);
return 0;
}
总结
好多树形DP都可以用线段树合并来优化啊……
