[JZOJ 5911] [NOIP2018模拟10.18] Travel 解题报告 （期望+树形DP）

题目链接：

http://172.16.0.132/senior/#contest/show/2530/1

题目：

        EZ同学家里非常富有，但又极其的谦虚，说话又好听，是个不可多得的人才。
        EZ常常在假期环游世界，他准备去N（N<=100000）个国家之多，一些国家有航线连接，由于EZ同学有一定的强迫症，任意两个国家之间都能通过航路直接或间接到达，并且这样的路径仅有一种。（简单来说，这些国家构成了一棵树）
         由于EZ是C国人，因此将C国（1号国家）作为整棵树的根
        每个国家有一个旅游热度A[i]和影响力D[i]。由于目的地有点多，为了避免选择困难症，他给每个国家设置了一个向往值F[i]，它等于所有的A[j]之和，满足i国在j国向C国走D[j]步的路径上（经过一条航路算一步，i=j也会被统计，如果D[j]步超过了C国，则超出部分不用管）。
         LYD同学家里有矿，富有程度与EZ不相上下，但他却在宅与现充间摇摆不定。某次机缘巧合，EZ外出旅游刺激了LYD，他决定也要开始旅游。为了避免又被判高重复率导致被取消资格，他将EZ的旅游地图略微做了一点调整，每条航路将有一定的概率出现。
         现在他有Q个询问，每次询问某个国家所在的联通块（由于每条边是一定概率出现，因此它所在的联通块可以是很多种）中所有国家的F[i]值的和的平方的期望（对998244353取模），以此来决定他旅游的目的地。但他极其厌恶繁琐的计算，于是他找到了能算出圆周率并将它倒背下来的你，答应给你丰厚的报酬。家里没矿，老爸也不是X达集团老总的你决定接受他的任务。

题外话：

这题...我查错差不多4小时吧，这一次查错发现了很多代码应该注意的坑点，我无一例外踩了一个遍

题解：

注意1：注意！！！注意！！！期望和的平方不等于和的平方的期望

注意2：每个在EZ的地图中是没有出现概率的说法的，因此每个国家的f值与边的出现概率无关

首先考虑怎么求f数组，对于每个国家在它自己打上一个+的标记，在它的 $d[i]+1$ 级祖先打上-的标记， 就可以直接子树求和了。查询某个位置的 $d[i]+1$ 级祖先可以在 DFS 的时候维护一 个栈，然后对于每个点直接访问栈中相应位置即可，这也是线性的。因此这样就能线性的求出$f$数组

考虑如何计算和的平方，我们不妨以询问的国家作为根重新 DFS 一遍这棵树，考虑 DP, 主要问题就是求两个块以一条出现概率为$ p$ 的边合并时的答案。不妨令原本的块期望和为$a$，并入的块期望和为$b$，根据期望的线性可加性，那么答案是$p(a + b)^2 + (1 − p)a^2 = a^2 + p(2ab +b ^2 )$

这样对于每个节点，维护$k2$数组表示以它为根的子树的期望和，维护$k3$数组表示以它为根的子树的和平方的期望，根据上面那个式子我们就可以合并答案

由于每次询问都要重新DFS，时间复杂度 $O(NQ)$，显然不行（事实上这样就30分）

其实并不需要对于每次询问都重新 DFS 转移，我们不妨仍然以1为根，那么 对于一个节点$i$，我们只需要知道$i$的子树部分的$k2$和$k3$,以及$i$以外的部分的期望和以及和平方的期望，这样就可以维护出i的任意一个儿子的上述信息，我们设$g[i]$表示除i的子树之外的节点的和平方的期望，h[i]表示除i的子树外的节点的期望和，$ans[i]=g[y]+2*h[y]*k2[y]+k3[y]$(注意不能写成$ans[i]=(h[y]+k2[y])^2$，原因参考注意1)

这就是二次扫描与换根的思想

具体怎么从i转移到$i$的儿子呢？我们给i的儿子一个顺序，对第$j$个儿子维护前缀其他儿子的两个信息，再维护后缀其他儿子的两个，这样第j个儿子的信息就相等于合并4部分，$h[x]$，前缀，合并，$f[x]$（代码中表示为$a,b,c,d$）

至于我查错中出现的坑点，一般仅限我的代码，就不赘述了

#include<algorithm>
#include<cstring>
#include<cstdio>
#include<iostream>
using namespace std;
typedef long long ll;

const int N=4e5+15;
const ll mo=998244353;
const int S=1e7;
int n,tot=1,tp;
int h[N],d[N],sta[N],tt[N];
ll cha[N],a[N],k2[N],k3[N],p1[N],p2[N],p3[N],f[N],H[N],G[N],ans[N];
ll Pre2[N],Pre3[N],Bac2[N],Bac3[N];
struct E{
    int to,nxt;ll p;
}e[N<<1];
inline int gc(){
    static char buf[S];
    static int len=0,pos=0;
    if (pos==len) pos=0,len=fread(buf,1,S,stdin);
    if (pos==len) exit(0);
    return buf[pos++];
}
inline ll read(){
    char ch=gc();ll s=0,f=1;
    while (ch<'0'||ch>'9') {if (ch=='-') f=-1;ch=gc();}
    while (ch>='0'&&ch<='9') {s=(s<<3)+(s<<1)+ch-'0';ch=gc();}
    return s*f;
}
void link(int u,int v,ll p) {e[++tot]=(E){v,h[u],p};h[u]=tot;}
void dfs1(int x,int pre){
    sta[++tp]=x;
    cha[x]=(cha[x]+a[x])%mo;
    cha[sta[tp-d[x]-1]]=(cha[sta[tp-d[x]-1]]+(mo-a[x]%mo))%mo;
    for (int i=h[x],y;i;i=e[i].nxt){
        if ((y=e[i].to)==pre) continue;
        dfs1(y,x);
    }
    tp--;
}
void dfs2(int x,int pre){
    f[x]=cha[x];
    for (int i=h[x],y;i;i=e[i].nxt){
        if ((y=e[i].to)==pre) continue;
        dfs2(y,x);
        f[x]=(f[x]+f[y])%mo;
    }
}
void dfs3(int x,int pre){
    k2[x]=f[x];
    k3[x]=f[x]*f[x]%mo;
    for (int i=h[x],y;i;i=e[i].nxt){
        if ((y=e[i].to)==pre) continue;
        ll p=e[i].p;
        dfs3(y,x);
        k3[x]=(k3[x]+2*p%mo*k2[x]%mo*k2[y]%mo+p*k3[y]%mo)%mo;
        k2[x]=((k2[x]+k2[y])*p%mo+((1-p)%mo+mo)%mo*k2[x]%mo)%mo;
    }
}
void dfs4(int x,int pre){
    Pre2[0]=0;Pre3[0]=0;
    int j=0;
    for (int i=h[x],y;i;i=e[i].nxt){
        if ((y=e[i].to)==pre) continue;
        ++j;
        tt[j]=i;
        ll p=e[i].p;
        Pre3[j]=(Pre3[j-1]+2*p%mo*Pre2[j-1]%mo*k2[y]%mo+p*k3[y]%mo)%mo;
        Pre2[j]=((Pre2[j-1]+k2[y])*p%mo+((1-p)%mo+mo)%mo*Pre2[j-1]%mo)%mo;
    }
    Bac2[j+1]=0;Bac3[j+1]=0;
    for (;j;j--){
        int i=tt[j];
        int y=e[i].to;
        ll p=e[i].p;
        Bac3[j]=(Bac3[j+1]+2*p%mo*Bac2[j+1]%mo*k2[y]%mo+p*k3[y]%mo)%mo;
        Bac2[j]=((Bac2[j+1]+k2[y])%mo*p%mo+((1-p)%mo+mo)%mo*Bac2[j+1]%mo)%mo;
        
        ll a=H[x],b=Pre2[j-1],c=Bac2[j+1],d=f[x];
        ll aa=G[x],bb=Pre3[j-1],cc=Bac3[j+1],dd=f[x]*f[x]%mo;
        G[y]=aa;//一个个合并，不能一起，参考注意1 
        H[y]=a;
        G[y]=(G[y]+bb+2*H[y]%mo*b%mo)%mo;
        H[y]=(H[y]+b)%mo;
        G[y]=(G[y]+cc+2*H[y]%mo*c%mo)%mo;
        H[y]=(H[y]+c)%mo;
        G[y]=(G[y]+dd+2*H[y]%mo*d%mo)%mo;
        H[y]=(H[y]+d)%mo;
        G[y]=G[y]*p%mo;H[y]=H[y]*p%mo;
        ans[y]=(G[y]+2*H[y]%mo*k2[y]%mo+k3[y])%mo;
    }
    for (int i=h[x],y;i;i=e[i].nxt){
        if ((y=e[i].to)==pre) continue;
        dfs4(y,x);
    }
}
int main()
{
    freopen("travel.in","r",stdin);
    freopen("travel.out","w",stdout);
    n=read();
    for (int i=1;i<=n;i++) a[i]=read(),d[i]=read();
    ll p;
    for (int i=1,u,v;i<n;i++){
        u=read();v=read();p=read();
        link(u,v,p);
        link(v,u,p);
    }
    dfs1(1,0);
    dfs2(1,0);
    dfs3(1,0);
    ans[1]=k3[1];
    dfs4(1,0);
    int q=read();
    while (q--){
        int s=read();
        printf("%lld\n",ans[s]);
    }
    return 0;
}