luogu P8334 [ZJOI2022] 深搜
首先显然先差分,然后枚举最小值,统计每个点对之间最小值至少为\(x\)的期望,如果我们将大于等于\(x\)的点设为黑点,小于等于\(x\)的点设为白点,两个点之间内通过黑点走到要满足在每一步走进目标点子树的时候不能走进有白色节点的子树。
可以考虑一个dp:设\(f_i\)表示\(i\)点走到子树内节点的期望,同时设有白色节点的子树个数为\(cnt\),则分类讨论:
如果要转移的子树内有白色节点,则\(f_i+=\frac{f_v}{cnt}\)。
如果要转移的子树内没有白色节点,则\(f_i+=\frac{f_v}{cnt+1}\)。
第一个转移可以看作是\(cnt\)个点随机排列,你指定的点在第一个的期望,显然是\(\frac{1}{cnt}\)。
第二个转移可以将这个白点看成黑点,那么就是和上面一个一样的结构。
答案为所有\(f_i\)之和。
也就是说我们现在要支持两个操作:将一个点改变颜色,将一个点子树内是否含有白色节点的状态改变。
容易发现只有\(O(n)\)次操作。于是可以直接ddp维护,设矩阵三个值为当前dp值,答案总和和常数。树剖的时候维护点东西即可。
时间复杂度\(O(n\log^2n3^3)\)需要卡常。
code:
#include<bits/stdc++.h>
#define Gc() getchar()
#define Me(x,y) memset(x,y,sizeof(x))
#define Mc(x,y) memcpy(x,y,sizeof(x))
#define d(x,y) ((m)*(x-1)+(y))
#define R(n) (rnd()%(n)+1)
#define Pc(x) putchar(x)
#define LB lower_bound
#define UB upper_bound
#define PB push_back
using ll=long long;using db=double;using lb=long db;using ui=unsigned;using ull=unsigned ll;
using namespace std;const int N=4e5+5,M=N*5+5,K=1e5+5,mod=998244353,Mod=mod-1;const db eps=1e-5;const int INF=1e9+7;
ll mpow(ll x,int y=mod-2){ll Ans=1;while(y) y&1&&(Ans=Ans*x%mod),y>>=1,x=x*x%mod;return Ans;}
int n,m,k,x,y,z,Rt,Df[N],A[N],B[N],d[N],Fl[N],vis[N],fa[N],Ct[N],Tp[N],En[N],Si[N],Sn[N],Id[N],Ih;ll ToT,Iv[N],g1[N],g2[N],h[N],frc[N];vector<int> S[N];
bool cmp(int x,int y){return A[x]<A[y];}
void dfs1(int x,int La){fa[x]=La;Si[x]=1;d[x]=d[La]+1;for(int i:S[x]) i^La&&(dfs1(i,x),Si[x]+=Si[i],Si[Sn[x]]<Si[i]&&(Sn[x]=i));}
void dfs2(int x,int La){Tp[x]=La;ToT+=d[x];Df[Id[x]=++Ih]=x;En[Tp[x]]=x;g1[x]=Si[x]-1-Si[Sn[x]];g2[x]=0;if(!Sn[x]) return;dfs2(Sn[x],La);h[x]=ToT;for(int i:S[x]) i^fa[x]&&i^Sn[x]&&(dfs2(i,i),0);h[x]=(ToT-h[x]-1ll*d[x]*g1[x])%mod;}
//void GA(int x,int La,int w){f[x]=0;P1[x]=Iv[Ct[x]+1];P2[x]=Iv[Ct[x]];for(int i:S[x]) i^La&&(GA(i,x,w),f[x]=(f[x]+f[i]*(vis[i]?P2[x]:P1[x]))%mod);if(Fl[x]) f[x]=0;else f[x]++,ToT+=w*f[x]%mod;}
struct MaT{ll f[3][3];MaT(){Me(f,0);}MaT operator *(const MaT &B)const{int i,j,h;MaT C;for(h=0;h<3;h++) for(i=0;i<3;i++) for(j=0;j<3;j++) C.f[i][j]+=f[i][h]*B.f[h][j];for(i=0;i<3;i++)for(j=0;j<3;j++) C.f[i][j]%=mod;return C;}}Cl,p;
MaT Fs[M],*Fh;
void GP(int x,MaT &B){B=Cl;B.f[2][1]=(h[x]%mod+mod)%mod;if(Fl[x])return;B.f[0][0]=Iv[Ct[x]+(!vis[Sn[x]])];B.f[2][0]=(g1[x]%mod*Iv[Ct[x]+1]+g2[x]%mod*Iv[Ct[x]]+1)%mod;}
struct Tree{
#define ls v<<1
#define rs v<<1|1
MaT *f;int Ns,Ba;void Up(int v){f[v]=f[rs]*f[ls];}void BD(int l,int r,int v=1){if(l==r) return GP(Df[l+Ba],f[v]);int m=l+r>>1;BD(l,m,ls);BD(m+1,r,rs);Up(v);}
void Ins(int x,int l,int r,int v=1){if(l==r) return GP(Df[l+Ba],f[v]);int m=l+r>>1;x<=m?Ins(x,l,m,ls):Ins(x,m+1,r,rs);Up(v);}
void Make(int x,int y){f=Fh;Fh+=4*x+1;Ns=x;Ba=y-1;BD(1,Ns);}
#undef ls
#undef rs
}Ro[N];
void Del(int x){MaT p;x=Tp[x];while(fa[x]) p=Ro[x].f[1],vis[x]?(g2[fa[x]]+=mod-p.f[2][0]):(g1[fa[x]]+=mod-p.f[2][0]),h[fa[x]]-=p.f[2][0]+p.f[2][1],x=Tp[fa[x]];}
void Ins(int x){MaT p;Ro[Tp[x]].Ins(Id[x]-Id[Tp[x]]+1,1,Ro[Tp[x]].Ns);
x^Tp[x]&&(Ro[Tp[x]].Ins(Id[fa[x]]-Id[Tp[x]]+1,1,Ro[Tp[x]].Ns),0);x=Tp[x];while(fa[x])
p=Ro[x].f[1],vis[x]?(g2[fa[x]]+=p.f[2][0]):(g1[fa[x]]+=p.f[2][0]),h[fa[x]]+=p.f[2][0]+p.f[2][1],Ro[Tp[fa[x]]].Ins(Id[fa[x]]-Id[Tp[fa[x]]]+1,1,Ro[Tp[fa[x]]].Ns),x=Tp[fa[x]];}
void Solve(){
int i,j;scanf("%d%d",&n,&Rt);for(i=1;i<=n;i++) scanf("%d",&A[i]),B[i]=i;sort(B+1,B+n+1,cmp);for(i=1;i<=n;i++) S[i].clear();for(i=1;i<n;i++) scanf("%d%d",&x,&y),S[x].PB(y),S[y].PB(x);
for(frc[0]=i=1;i<=n+1;i++) frc[i]=frc[i-1]*i%mod;Iv[n]=mpow(frc[n]);for(i=n-1;~i;i--) Iv[i]=Iv[i+1]*(i+1)%mod;for(i=1;i<=n;i++) Iv[i]=Iv[i]*frc[i-1]%mod;
Fh=Fs;Me(h,0);Me(Fl,0);Me(vis,0);Me(Ct,0);ToT=0;Me(Sn,0);Ih=0;dfs1(Rt,0);dfs2(Rt,Rt);for(i=1;i<=n;i++) Tp[i]==i&&(Ro[i].Make(d[En[i]]-d[i]+1,Id[i]),0);
ToT=ToT*A[B[1]]%mod;for(i=1;i<n;i++) {x=B[i];Del(x);Fl[x]=1;Ins(x);
while(fa[x]&&!vis[x]) Del(x),vis[x]=1,Ct[fa[x]]++,Ins(x),x=fa[x];p=Ro[Rt].f[1];
ToT+=(A[B[i+1]]-A[B[i]])*(p.f[2][1]+p.f[2][0])%mod;}printf("%lld\n",ToT%mod);
}
int main(){
freopen("1.in","r",stdin);
int T;Cl.f[0][1]=Cl.f[1][1]=Cl.f[2][2]=1;scanf("%d",&T);while(T--) Solve();
}