【洛谷6773】[NOI2020] 命運（線段樹合併優化DP）

阿新 • • 發佈：2020-12-17

點此看題面

給定一棵$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;//在根節點的線段樹上詢問答案
}