期望+树形dp--xjoi 2018提高组模拟赛13 T2

F数组通过差分求，这里学到了一个小技巧，dfs时维护一个栈可以在访问某点时O（1）查询k级祖先。之后和的平方的期望要拆开这个平方，发现其实就是所有有序数对相乘期望，这样可以考虑树形dp时合并，维护期望和f,和的平方期望g数组，可知道两块合并（假设2号并到1号上）,g1=p(g1+g2+2f1f2)+(1-p)g1   f1=p(f1+f2)+(1-p)f1.这里p指连接两块的边出现概率。不过如果对于每个询问都去dfs一遍复杂度就炸了，所以要用到一个小技巧，先以1为根算出1答案，同时在此时所有父亲都计算了各自孩子的贡献，之后再dfs一波，从1开始，考虑当前父亲节点更新孩子，因为父亲的f，g数组都已考虑整棵树，所以先去掉被当前要更新的孩子所更新的贡献，再用这个东西去更新孩子即可。（简单说就是一遍孩子更新父亲，再一遍父亲更新孩子就完了。。）

代码：

#include<bits/stdc++.h>
#define ll long long
using namespace std;
const int N=2e5+10,M=N*2;
const ll mod=998244353;
 
int n,d[N];
int hd[N],nxt[M],to[M],w[M],tot;
int stk[N],snum;
ll a[N],f[N],sum[N],spf[N];
 
void Ad(ll &x,ll y)
{x=(x+y)%mod;}
void Dw(ll &x,ll y)
{x=(x-y+mod)%mod;}
void init()
{
    memset(hd,-1,sizeof hd);
    tot=-1;
}
void add(int u,int v,int ww)
{
    nxt[++tot]=hd[u],to[tot]=v,w[tot]=ww,hd[u]=tot;
    nxt[++tot]=hd[v],to[tot]=u,w[tot]=ww,hd[v]=tot;
}
void get_f(int pos,int fa,int dep)
{
    Ad(f[pos],a[pos]),stk[++snum]=pos;
    Dw(f[stk[max(0,dep-d[pos]-1)]],a[pos]);
    for(int i=hd[pos];i!=-1;i=nxt[i])
    {
        if(to[i]==fa)continue;
        get_f(to[i],pos,dep+1),Ad(f[pos],f[to[i]]);
    }
    --snum;
}
void dfs1(int pos,int fa)
{
    int v;
    sum[pos]=f[pos],spf[pos]=f[pos]*f[pos]%mod;
    for(int i=hd[pos];i!=-1;i=nxt[i])
    {
        v=to[i];
        if(to[i]==fa)continue;
        dfs1(v,pos);
        spf[pos]=((spf[pos]+spf[v]+2LL*sum[pos]*sum[v]%mod)%mod*w[i]%mod+(1LL-w[i]+mod)%mod*spf[pos]%mod)%mod;
        sum[pos]=((sum[pos]+sum[v])%mod*w[i]%mod+(1LL-w[i]+mod)*sum[pos]%mod)%mod;
    }
}
void dfs2(int pos,int fa)
{
    int v;ll F,S;
    for(int i=hd[pos];i!=-1;i=nxt[i])
    {
        v=to[i];
        if(v==fa)continue;
        S=spf[pos],F=sum[pos];
        Dw(F,w[i]*sum[v]%mod);
        Dw(S,w[i]*spf[v]%mod);
        Dw(S,2LL*w[i]*F%mod*sum[v]%mod);
        spf[v]=((spf[v]+S+2LL*sum[v]*F%mod)%mod*w[i]%mod+(1LL-w[i]+mod)%mod*spf[v]%mod)%mod;
        sum[v]=((sum[v]+F)%mod*w[i]%mod+(1LL-w[i]+mod)*sum[v]%mod)%mod;
        dfs2(v,pos);
    }
}
int main()
{
    int u,v,ww,q;
    init();
    scanf("%d",&n);
    for(int i=1;i<=n;i++)
        scanf("%lld%d",&a[i],&d[i]);
    for(int i=1;i<n;i++)
        scanf("%d%d%d",&u,&v,&ww),add(u,v,ww);
    get_f(1,0,1),dfs1(1,0),dfs2(1,0);
    scanf("%d",&q);
    while(q--)
    {
        scanf("%d",&u);
        printf("%lld\n",spf[u]);
    }
}