【洛谷6773】[NOI2020] 命运（线段树合并优化DP）

点此看题面

• 给定一棵\(n\)个点的树，每条边可以填上\(0\)或\(1\)。
• 给出\(m\)个限制，每个表示树上一条从上向下的直链中至少有一条边为\(1\)。
• 求有多少种合法的方案。
• \(n,m\le5\times10^5\)

---

不得不说，现在看来这道题真的挺水的。

主要线段树合并优化\(DP\)这个套路我早就接触过了（【洛谷5298】[PKUWC2018] Minimax），但这里还是稍微再简单提一下吧。

暴力\(DP\)

我们设\(f_{x,i}\)表示\(x\)子树内所有边权都已经确定时，尚未满足的限制中最大的深度为\(i\)的方案数。

考虑从子节点\(y\)向\(x\)的转移，无非就是\(x\)与\(y\)之间的边填\(0\)还是\(1\)两种情况：

• 填\(0\)：则比较\(x,y\)原本的\(i\)的大小，分两类转移，有\(f_{x,i}=\sum_{j=0}^if_{x,i}\times f_{y,j}+\sum_{j=0}^{i-1}f_{x,j}\times f_{y,i}\)。
• 填\(1\)：那么所有限制都能得到满足，有\(f_{x,i}=\sum_{j=0}^{dep_x}f_{x,i}\times f_{y,j}\)。

线段树合并优化

令\(S_{x,i}=\sum_{j=0}^if_{x,i}\)，然后把上面的式子稍微理一理，得到：

\[f_{x,i}=f_{x,i}\times(S_{y,dep_x}+S_{y,i})+f_{y,i}\times S_{x,i-1} \]

发现其中的\(S_{y,dep_x}\)是个定值，可以在转移前先询问得出。

而其余的项都是下标与\(i\)有关的前缀和，只要在线段树合并的时候，维护好当前区间左侧的前缀和（\(S_{x,l-1}\)和\(S_{y,l-1}\)）就好了。

具体实现还是详见代码吧。

代码：\(O(nlogn)\)

#include<bits/stdc++.h>
#define Tp template<typename Ty>
#define Ts template<typename Ty,typename... Ar>
#define Reg register
#define RI Reg int
#define Con const
#define CI Con int&
#define I inline
#define W while
#define N 500000
#define X 998244353
#define add(x,y) (e[++ee].nxt=lnk[x],e[lnk[x]=ee].to=y)
using namespace std;
int n,ee,lnk[N+5];struct edge {int to,nxt;}e[N<<1];
class SegmentTree
{
	private:
		#define PT CI l=0,CI r=n
		#define LT l,mid
		#define RT mid+1,r
		#define PU(x) (O[x].V=(O[O[x].S[0]].V+O[O[x].S[1]].V)%X)
		#define T(x,v) x&&(O[x].V=1LL*O[x].V*v%X,O[x].F=1LL*O[x].F*v%X)
		#define PD(x) O[x].F^1&&(T(O[x].S[0],O[x].F),T(O[x].S[1],O[x].F),O[x].F=1)
		int Nt;struct node {int V,F,S[2];}O[N<<6];
	public:
		I void U(int& rt,CI x,PT)//单点修改
		{
			if(!rt&&(O[rt=++Nt].F=1),++O[rt].V,l==r) return;RI mid=l+r>>1;PD(rt);
			x<=mid?U(O[rt].S[0],x,LT):U(O[rt].S[1],x,RT);
		}
		I int Q(CI rt,CI L,CI R,PT)//区间查询
		{
			if(!rt) return 0;if(L<=l&&r<=R) return O[rt].V;RI mid=l+r>>1;PD(rt);
			return ((L<=mid?Q(O[rt].S[0],L,R,LT):0)+(R>mid?Q(O[rt].S[1],L,R,RT):0))%X;
		}
		I void Merge(int& x,CI y,CI s1,CI s2,PT)//线段树合并，s1和s2维护区间左侧前缀和
		{
			if(!x||!y) return (void)(x&&T(x,s1),y&&(x=y,T(x,s2)));RI mid=l+r>>1;PD(x),PD(y);//只有一个点
			if(l==r) return (void)(O[x].V=(1LL*O[x].V*(s1+O[y].V)+1LL*O[y].V*s2)%X);//叶节点
			Merge(O[x].S[1],O[y].S[1],(s1+O[O[y].S[0]].V)%X,(s2+O[O[x].S[0]].V)%X,RT),//先做右区间
			Merge(O[x].S[0],O[y].S[0],s1,s2,LT),PU(x);//再做左区间，防止修改左区间影响右区间的转移
		}
}S;
int dep[N+5];I void Init(CI x=1,CI lst=0,CI d=1)//初始化，求出所有点深度
{
	dep[x]=d;for(RI i=lnk[x];i;i=e[i].nxt) e[i].to^lst&&(Init(e[i].to,x,d+1),0);
}
int v[N+5],Rt[N+5];I void dfs(CI x=1,CI lst=0)//dfs利用线段树合并做一遍DP
{
	S.U(Rt[x],v[x]);for(RI i=lnk[x];i;i=e[i].nxt) e[i].to^lst&&//DP初始化当前点的限制点中深度最大的点
		(dfs(e[i].to,x),S.Merge(Rt[x],Rt[e[i].to],S.Q(Rt[e[i].to],0,dep[x]),0),0);//从子节点上传信息
}
int main()
{
	RI i,x,y;for(scanf("%d",&n),i=1;i^n;++i) scanf("%d%d",&x,&y),add(x,y),add(y,x);
	RI Qt;Init(),scanf("%d",&Qt);W(Qt--) scanf("%d%d",&x,&y),v[y]=max(v[y],dep[x]);//维护深度最大的点
	return dfs(),printf("%d\n",S.Q(Rt[1],0,0)),0;//在根节点的线段树上询问答案
}