洛谷P3412 仓鼠找sugar II 题解

题意：

小仓鼠和他的基（mei）友（zi）sugar 住在地下洞穴中，每个节点的编号为 \(1\sim n\) 间的一个整数。地下洞穴是一个树形结构（两个洞穴间距离为 \(1\)）。这一天小仓鼠打算从从他的卧室 \(a\)（是任意的）走到他的基友卧室 \(b\)（还是任意的）（注意，\(a\) 有可能等于 \(b\)）。然而小仓鼠学 OI 学傻了，不知道怎么怎么样才能用最短的路程走到目的地，于是他只能随便乱走。当他在一个节点时，会等概率走到这个点的母亲或者所有孩子节点（例如这个节点有一个母亲节点和两个子节点，那么下一步走到这 \(3\) 个节点的概率都是 \(\dfrac{1}{3}\)），一旦走到了他基友的卧室，就会停下。

现在小仓鼠希望知道，他走到目的地时，走的步数的期望。这个期望本来是一个有理数，但是为了避免误差，我们要求输出这个有理数对 \(998,422,353\) 取模的结果。

形式化地说，可以证明答案可以被表示为既约分数 \(\dfrac{y}{x}\)，其中 \(x\not\equiv 0\pmod {998,422,353}\)。可以证明存在一个唯一的整数 \(z\)（\(0\le z\lt 998,422,353\)），使得 \(xz=y\)，你只需要输出 \(z\) 的值。

小仓鼠那么弱，还要天天被 JOHNKRAM 大爷虐，请你快来救救他吧！

输入格式：

第一行一个正整数 \(n\)，表示这棵树节点的个数。

接下来 \(n-1\) 行，每行两个正整数 \(u\) 和 \(v\)，表示节点 \(u\) 到节点 \(v\) 之间有一条边。

输出格式：

一个整数，表示取模后的答案。

数据范围：

\(1\le n\le 100000\)。

考虑dp。设 \(g_u\) 表示从 \(u\) 走到 \(fa(u)\) 的期望步数，其中 \(u \ne \mathrm{root}\) ，则 \[ g_u = \frac{1}{d_u}\left(1 + \sum_{fa(v)=u}\left(g_v + g_u + 1\right)\right) \] 其中 \(d_u\) 表示点 \(u\) 的度数。

解一下方程 \[ g_u = d_u + \sum_{fa(v) = u} g_v \] 设 \(f_u\) 表示 \(u\) 子树内所有点全部走到点 \(u\) 的期望步数，则 \[ f_u = \sum_{fa(v) = u} \left(f_v + \mathrm{size}(v)\cdot g_v\right) \] 考虑换根dp。设 \(h_u\) 表示 \(u\) 为根时的 \(f_u\) 值，则 \[ h_u = f_u + (h_{fa(u)} - f_u - \mathrm{size}(u)\cdot g_u) + (n-\mathrm{size}(u))\cdot(g_1-g_u) \] 则答案为 \[ \mathrm{Ans} = \frac{\sum_{i=1}^n h_i}{n^2} \] 时间复杂度 \(\mathcal{O}(n)\)

代码：

#include <bits/stdc++.h>
using namespace std;
#define int long long
#define INF 0x3f3f3f3f3f3f3f3f
const int mod = 998244353;
void up(int &x,int y) { x < y ? x = y : 0; }
void down(int &x,int y) { x > y ? x = y : 0; }
void add(int &x, int y) { (x += y) >= mod ? x -= mod : 0; }
#define N ((int)(1e5 + 15))

int n, inv, res, in[N], f[N], g[N], h[N], sz[N];
vector<int> vec[N];
int qpow(int a, int b)
{
    int r = 1;
    while(b)
    {
        if(b & 1) r = r * a % mod;
        b >>= 1; a = a * a % mod;
    }
    return r;
}
void dfs1(int u, int fa)
{
    g[u] = in[u]; sz[u] = 1;
    for(int v : vec[u]) if(v != fa)
    {
        dfs1(v, u); sz[u] += sz[v];
        add(g[u], g[v]); add(f[u], (f[v] + sz[v] * g[v]) % mod);
    }
}
void dfs2(int u, int fa)
{
    h[u] = (f[u] + (h[fa] - f[u] + mod - sz[u] * g[u] % mod + mod)) % mod;
    add(h[u], (n - sz[u]) * (g[1] - g[u] + mod) % mod); add(res, h[u]);
    for(int v : vec[u]) if(v != fa) { dfs2(v, u); }
}
signed main()
{
    ios::sync_with_stdio(0);
    cin.tie(0);cout.tie(0);
    // freopen("check.in","r",stdin);
    // freopen("check.out","w",stdout);
    cin >> n; inv = qpow(n, mod - 2);
    for(int i = 1, u, v; i < n; i++)
    {
        cin >> u >> v;
        vec[u].push_back(v); vec[v].push_back(u); ++in[u]; ++in[v];
    }
    dfs1(1, 0); res = f[1]; h[1] = f[1];
    for(int v : vec[1]) dfs2(v, 1);
    res = res * inv % mod * inv % mod;
    cout << res << '\n';
    return 0;
}