[JZOJ5519] Hello my friend

Description

给出一棵N个节点的树，每个点有可能是黑白两种颜色的一种

现在从1号点开始随机游走（即走这个点的每条出边的概率是相同的），每到一个点，如果这个点是黑点，或者这是白点并且这个点第一次经过，那么答案+1。当走到度数为1的节点时游走停止（保证1号节点的度数大于1）

求答案的期望（对998244353取模）
N<=100000

Solution

不妨先考虑只有黑点怎么做

设$F_i$表示i号点到结束的期望答案

明显F[叶子]=1

可以从下到上树形DP
列出F的式子
令j为i的儿子,$d_i$为i的度数（包括父亲），x为i的父亲
$F_i=1+{F_x+\sum\limits F_j\over d_i}$

不妨假设儿子的已经求出来了，是一个常数，那么每个点的F可以表示成$F_i=k_iF_x+b_i$
$F_i=1+{F_x+\sum k_jF_i+\sum b_j\over d_i}$
化简得
$(d_i-\sum k_j)F_i=d_i+F_x+\sum b_j$
$F_i={F_x\over d_i-\sum k_j}+{d_i+\sum b_j\over d_i-\sum k_j}$

这就表示成了我们想要的形式

容易得出k[叶子]=0,b[叶子]=1
那么直接树形DP上去即可
答案就是$b_1$（因为$k_1=0$）

考虑只有白点怎么做

令$h_i$表示i能走到x的概率（x是i的父亲）

那么$h_i={1+\sum h_j*h_i\over d_i}$
化简得$h_i={1\over d_i-\sum h_j}$

从下到上树形DP

令$g_i$表示x能走到i的概率，p为x的儿子
$g_i={1+g_x*g_i+\sum\limits_{p\not=i} h_p*g_i\over d_x}$

化简得$g_i={1\over d_x-\sum\limits_{p\not=i} h_p-g_x}$

从上到下树形DP
那么每个白点的贡献就是1走到这个点的概率
g一路乘下来就是了

考虑黑白点都有

结合黑点的算法和白点的算法，白点单独计算，计算黑点的时候如果走到白点就把那个1去掉就行了

Code

#include <cstdio>
#include <algorithm>
#include <iostream>
#include <cmath>
#include <cstdlib>
#include <cstring>
#define fo(i,a,b) for(int i=a;i<=b;i++)
#define fod(i,a,b) for(int i=a;i>=b;i--)
#define mo 998244353
#define N 100005
#define LL long long
using namespace std;
int n,m,nt[2*N],fs[N],dt[2*N],c[N],rd[N];
LL f[N],b[N],g[N],h[N],ans,hs[N];
void link(int x,int y)
{
    nt[++m]=fs[x];
    dt[fs[x]=m]=y;
}
LL ksm(LL k,LL n)
{
    LL s=1;
    for(;n;k=k*k%mo,n>>=1) if(n&1) s=s*k%mo;
    return s;
}
void dp(int k,int fa)
{
    if(nt[fs[k]]==0)
    {
        f[k]=0,b[k]=c[k];
        h[k]=0;
        return;
    }
    LL sk=0,bs=0;
    hs[k]=0;
    for(int i=fs[k];i;i=nt[i])
    {
        int p=dt[i];
        if(p!=fa)
        {
            dp(p,k);
            (sk+=f[p])%=mo,(bs+=b[p])%=mo,(hs[k]+=h[p])%=mo;
        }
    }
    LL ny=ksm((rd[k]-sk+mo)%mo,mo-2);
    h[k]=ksm((rd[k]-hs[k]+mo)%mo,mo-2);
    f[k]=ny;
    b[k]=((c[k]==1)?((rd[k]+bs)%mo*ny%mo):(bs*ny%mo));
}
void dfs(int k,int fa,LL v)
{
    if(k==1) g[k]=0;
    else g[k]=ksm(((rd[fa]-hs[fa]+h[k]-g[fa])%mo+mo)%mo,mo-2);
    if(k!=1) v=v*g[k]%mo;
    if(c[k]==0) ans+=v;
    for(int i=fs[k];i;i=nt[i])
    {
        int p=dt[i];
        if(p!=fa) dfs(p,k,v);
    }
}
int main()
{
    freopen("0.in","r",stdin);
    freopen("0.out","w",stdout);
    scanf("%d\n",&n);
    fo(i,1,n)
    {
        char ch=getchar();
        c[i]=(ch=='1');
    }
    fo(i,1,n-1)
    {
        int x,y;
        scanf("%d%d",&x,&y);
        link(x,y),link(y,x),rd[x]++,rd[y]++;
    }
    dp(1,0);
    dfs(1,0,1);
    printf("%lld\n",(b[1]+ans)%mo);
}