[ZJOI2022] 深搜——树上动态DP、树链剖分

[ZJOI2022] 深搜

题解

其实这题虽然代码较复杂，但是思路比较简短。

按照惯例，既然是权值最小的那个，那么按点权从小到大（权值相同的任意定个顺序）考虑每个点的贡献。此时要求路径不能经过权值更小的点，并且能够经过当前节点。我们发现这两个要求非常苛刻，但是仅满足第一个要求貌似比较可做，所以我们转而求答案≥某个点权的概率之和，然后差分回来。

首先考虑树上有一些点不能经过时的答案。不妨把不能经过的点叫做黑点，设 $dp[x]$ 表示当 $y$ 为 $x$ 子树内的点时，从 $x$ 往 $y$ 深搜不经过黑点的概率之和。分类讨论一下转移：

• $x$ 是黑点，那么 $dp[x]=0$；
• $x$ 不是黑点，首先 $y=x$ 时有个1，然后考虑 $y$ 在 $x$ 的某个子树 $v$ 中，那么我们要求在深搜进入子树 $v$ 之前不能经过黑点。由于进入了其它子树内必然会遍历完整个子树，所以我们可以再定义“灰点”表示子树内有黑点的节点（黑点也是灰点），定义 $s{x}_x$ 表示 $x$ 的灰儿子数量，那么这部分的转移：
  $$dp[x]=1+\sum _{f{a}_v=x}\frac{dp[v]}{s{x}_x+[visnotgrey]}$$简要解释一下：如果 $v$ 是灰点，那么要求进入 $v$ 子树前不能进入其它灰子树。由于 $s{x}_x$ 个灰子树地位相等，所以 $v$ 成为 $s{x}_x$ 个子树中第一个的概率为 $\frac{1}{s{x}_x}$。$v$ 不是灰点同理。

这个转移就是个线性变换，所以我们用矩阵来表示它，每次添加黑点用动态DP维护即可。由于一个节点只会由白变灰，所以跳父亲修改 $sx$ 值是 $O(n)$ 的。

这题很卡，不能用一般的树链剖分，要用全局平衡，复杂度为 $O(n\log n)$。

代码

#include<bits/stdc++.h>//JZM yyds!!
#define ll long long//God JZM!!
#define lll __int128//JZM RollInDark!!
#define uns unsigned
#define fi first
#define se second
#define IF (it->fi)
#define IS (it->se)
#define lowbit(x) ((x)&-(x))
#define END putchar('\n')
#define inline jzmyyds
using namespace std;
const int MAXN=4e5+5;
const ll INF=1e18;
ll read(){
	ll x=0;bool f=1;char s=getchar();
	while((s<'0'||s>'9')&&s>0){if(s=='-')f^=1;s=getchar();}
	while(s>='0'&&s<='9')x=(x<<1)+(x<<3)+(s^48),s=getchar();
	return f?x:-x;
}
int ptf[50],lpt;
void print(ll x,char c='\n'){
	if(x<0)putchar('-'),x=-x;
	ptf[lpt=1]=x%10;
	while(x>9)x/=10,ptf[++lpt]=x%10;
	while(lpt)putchar(ptf[lpt--]^48);
	if(c>0)putchar(c);
}
const ll MOD=998244353;
ll ksm(ll a,ll b,ll mo){
	ll res=1;
	for(;b;b>>=1,a=a*a%mo)if(b&1)res=res*a%mo;
	return res;
}
struct matrix{
	int a,b,c,d;matrix(){}
	matrix(int A,int B,int C,int D){a=A,b=B,c=C,d=D;}
	matrix operator+(matrix y)const{
		y.a+=a,y.b+=b,y.c+=c,y.d+=d;
		if(y.a>=MOD)y.a-=MOD;
		if(y.b>=MOD)y.b-=MOD;
		if(y.c>=MOD)y.c-=MOD;
		if(y.d>=MOD)y.d-=MOD;
		return y;
	}
	matrix operator*(const matrix&y)const{
		return matrix((1ll*a*y.a+1ll*b*y.c)%MOD,
		(1ll*a*y.b+1ll*b*y.d)%MOD,
		(1ll*c*y.a+1ll*d*y.c)%MOD,
		(1ll*c*y.b+1ll*d*y.d)%MOD);
	}
};
int n,root,sr[MAXN];
ll inv[MAXN],a[MAXN],f[MAXN];
int init(int n){
	inv[0]=0,inv[1]=1;
	for(int i=2;i<=n;i++)inv[i]=(MOD-inv[MOD%i])*(MOD/i)%MOD;
	return 114514;
}
int cbddl=init(MAXN-4);
struct edge{
	int v,to;edge(){}
	edge(int V,int T){v=V,to=T;}
}e[MAXN<<1];
int EN,G[MAXN];
void addedge(int u,int v){
	e[++EN]=edge(v,G[u]),G[u]=EN;
	e[++EN]=edge(u,G[v]),G[v]=EN;
}
int fa[MAXN],siz[MAXN],hs[MAXN];
int hd[MAXN],id[MAXN],tl[MAXN],IN,tp[MAXN],rt[MAXN];
int sx[MAXN];
bool bp[MAXN];
ll sf[MAXN],sg[MAXN],ps[MAXN],pf[MAXN],ans;

struct itn{
	int ls,rs;
	matrix a,f,s;itn(){}
	itn(matrix A){s=f=a=A,ls=rs=0;}
}t[MAXN<<1];
void upd(int x){
	t[x].f=t[t[x].rs].f*t[x].a,t[x].s=t[t[x].rs].s+t[x].f;
	t[x].s=t[x].s+t[x].f*t[t[x].ls].s,t[x].f=t[x].f*t[t[x].ls].f;
}
matrix gtf(int x){
	if(bp[x])return matrix(1,0,0,0);
	if(bp[hs[x]]||sx[hs[x]]>0)
		return matrix(1,(1+sf[x]*inv[sx[x]+1]+sg[x]*inv[sx[x]])%MOD,0,inv[sx[x]]);
	return matrix(1,(1+sf[x]*inv[sx[x]+1]+sg[x]*inv[sx[x]])%MOD,0,inv[sx[x]+1]);
}
int build(int l,int r){
	if(l>r)return 0;
	int s=siz[id[l]]-siz[hs[id[r]]],p=0,x;
	for(x=l;x<r;x++){
		p+=siz[id[x]]-siz[hs[id[x]]];
		if((p<<1)>s)break;
	}t[x]=itn(gtf(id[x]));
	t[x].ls=build(l,x-1),t[x].rs=build(x+1,r);
	return upd(x),x;
}
void pushu(int x,int z){
	if(x==z){upd(x);return;}
	if(x>z)pushu(t[x].ls,z);
	else pushu(t[x].rs,z);
	upd(x);
}

void dfs1(int x){
	siz[x]=1,hs[x]=0;
	for(int i=G[x];i;i=e[i].to){
		int v=e[i].v;
		if(v==fa[x])continue;
		fa[v]=x,dfs1(v),siz[x]+=siz[v];
		if(siz[v]>siz[hs[x]])hs[x]=v;
	}
}
void dfs2(int x){
	hd[x]=++IN,id[IN]=x,bp[x]=0,sx[x]=0;
	tp[x]=(x==hs[fa[x]]?tp[fa[x]]:x),tl[tp[x]]=IN;
	sf[x]=siz[x]-1-siz[hs[x]],sg[x]=0;
	if(hs[x])dfs2(hs[x]);
	for(int i=G[x];i;i=e[i].to){
		int v=e[i].v;
		if(v==fa[x]||v==hs[x])continue;
		dfs2(v);
	}if(x==tp[x]){
		rt[x]=build(hd[x],tl[x]);
		pf[x]=siz[x],ps[x]=t[rt[x]].s.b,(ans+=ps[x])%=MOD;
	}
}
void update(int x){
	while(x){
		t[hd[x]].a=gtf(x),pushu(rt[tp[x]],hd[x]);
		x=tp[x],(ans+=MOD-ps[x])%=MOD;
		if(fa[x]){
			if(bp[x]||sx[x]>0)(sg[fa[x]]+=MOD-pf[x])%=MOD;
			else (sf[fa[x]]+=MOD-pf[x])%=MOD;
		}
		pf[x]=t[rt[x]].f.b,ps[x]=t[rt[x]].s.b,(ans+=ps[x])%=MOD;
		if(fa[x]){
			if(bp[x]||sx[x]>0)(sg[fa[x]]+=pf[x])%=MOD;
			else (sf[fa[x]]+=pf[x])%=MOD;
		}x=fa[x];
	}
}
int main()
{
 	freopen("dfs.in","r",stdin);
 	freopen("dfs.out","w",stdout);
	t[0].f=matrix(1,0,0,1);
	for(int NND=read();NND--;){
		n=read(),root=read(),fa[root]=0,IN=EN=ans=0;
		for(int i=1;i<=n;i++)a[i]=read(),sr[i]=i,G[i]=0;
		sort(sr+1,sr+1+n,[](int x,int y){return a[x]<a[y];});
		for(int i=1;i<n;i++)addedge(read(),read());
		dfs1(root),dfs2(root);
		for(int id=1;id<=n;id++){
			f[id]=ans;
			int x=sr[id];
			vector<int>c;
			c.push_back(x);
			while(fa[x]&&!bp[x]&&!sx[x])c.push_back(x=fa[x]);
			for(int x:c)if(x==tp[x]&&fa[x]){
				if(bp[x]||sx[x]>0)(sg[fa[x]]+=MOD-pf[x])%=MOD;
				else (sf[fa[x]]+=MOD-pf[x])%=MOD;
			}x=sr[id];
			for(int u:c){
				if(u==x)bp[u]=1;
				else sx[u]++;
				if(u==tp[u]&&fa[u]){
					if(bp[u]||sx[u]>0)(sg[fa[u]]+=pf[u])%=MOD;
					else (sf[fa[u]]+=pf[u])%=MOD;
				}
			}
			for(int u:c)update(u);
		}f[n+1]=0,ans=0;
		for(int i=n;i>0;i--)(ans+=(f[i]-f[i+1]+MOD)*a[sr[i]])%=MOD;
		print(ans);
	}
	return 0;
}