洛谷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;
}