【JZOJ5911】Travel

Description

有一棵以1为根的树，每个点有\(a_i\)，\(d_i\)，定义\(F_i\)为所有的\(a_j\)的和，满足\(j\)往根走\(d_j\)步的路径上有\(i\)。每条边有一定概率出现（对\(F\)无影响），Q次询问，每次询问\(x\)点所在的联通块所有\(F_i\)和的平方的期望。

Solution

首先求\(F_i\)，可以再树上查分，子树求和。
连通块很难处理，考虑一个点与它子树一部分组成连通块的期望，和很好处理，考虑和的平方，设当前块为\(a\)，要加入的子树块为\(b\) ，则新的块\(c=p(a+b)^2+(1-p)a^2\)（\(p\)为连接两个块的边的出现概率）
于是拆式子：\(c=a^2+p(2ab+b^2)\)，然后直接转移，根节点的答案即为询问该点的期望。
这题有多组询问，我们可以对\(1\)做完\(dp\)后，从上至下换根即可。

Code

#include<cstdio>
#include<cstdlib>
#include<cstring>
#include<algorithm>
#define fo(i,j,k) for(int i=j;i<=k;++i)
#define fd(i,j,k) for(int i=j;i>=k;--i)
#define rep(i,x) for(int i=ls[x];i;i=nx[i])
using namespace std;
typedef long long ll;
const int N=2e5+10,M=4e5+10,mo=998244353;
int a[N],d[N];
int to[M],nx[M],ls[N],vl[M],num=0;
ll F[N];
ll h[N],g[N];
void link(int u,int v,int w){
	to[++num]=v,nx[num]=ls[u],ls[u]=num;
	vl[num]=w;
}
int st[N],top=0;
void pre(int x,int fr){
	int fd=st[max(top-d[x],0)];
	st[++top]=x,F[x]+=a[x],F[fd]-=a[x];
	rep(i,x){
		int v=to[i];
		if(v==fr) continue;
		pre(v,x);
		F[x]+=F[v];
	}
	F[x]=(F[x]%mo+mo)%mo,st[top--]=0;
}
void dfs(int x,int fr){
	g[x]=F[x],h[x]=F[x]*F[x]%mo;
	rep(i,x){
		int v=to[i];
		ll p=vl[i];
		if(v==fr) continue;
		dfs(v,x);
		h[x]=(h[x]+(g[x]*g[v]*2+h[v])%mo*p%mo)%mo;
		g[x]=(g[x]+g[v]*p)%mo;
	}
}
ll G[N],H[N];
void calc(int x,int fr){
	rep(i,x){
		int v=to[i];
		ll p=vl[i];
		if(v==fr) continue;
		ll gv=(G[x]-g[v]*p%mo+mo)%mo,hv=(H[x]-(gv*g[v]%mo*2%mo+h[v])%mo*p%mo+mo)%mo;
		H[v]=(h[v]+(g[v]*gv%mo*2+hv)%mo*p%mo)%mo;
		G[v]=(g[v]+gv*p)%mo;
		calc(v,x);
	}
}
int main()
{
	freopen("travel.in","r",stdin);
	freopen("travel.out","w",stdout);
	int n;
	scanf("%d",&n);
	fo(i,1,n) scanf("%d %d",&a[i],&d[i]);
	fo(i,2,n){
		int u,v,w;
		scanf("%d %d %d",&u,&v,&w);
		link(u,v,w),link(v,u,w);
	}
	pre(1,0);
	dfs(1,0);
	G[1]=g[1],H[1]=h[1];
	calc(1,0);
	int q;
	scanf("%d",&q);
	while(q--){
		int x;
		scanf("%d",&x);
		printf("%lld\n",H[x]);
	}
}