S 老师的虚树 - 题解

标签与难度

标签: 期望, 贡献法, 树形DP, 组合数学, NTT, 多项式乘法难度: 2500

题目大意喵~

哈喵~！各位算法大师们，S 老师又来出题啦！这次是在一棵有 $n$ 个节点、 $n-1$ 条边的树上玩耍哦。每一条边都有自己的颜色，喵。

题目定义了一个叫做虚树的概念。对于一个点的集合 $S$ ，它的虚树是包含 $S$ 中所有点的、最小的连通子图。然后，这个虚树的权值 $w(S)$ 被定义为它所包含的边的不同颜色数量。

我们的任务是，对于每个 $i$ (从 1 到 $n$ )，计算当我们在 $n$ 个点中随机选择一个大小为 $i$ 的点集 $S$ 时，它的虚树权值 $w(S)$ 的期望是多少。最后结果要对 $998244353$ 取模，的说。

简单来说，对每个 $i \in [1, n]$ ，我们要计算：

E_i = \frac{\sum_{S \subseteq V, |S|=i} w(S)}{\binom{n}{i}}

解题思路分析

这道题看起来好复杂呀，又是虚树又是期望的，但别怕，让我带你一步步解开它的神秘面纱，喵~

第一步：期望的线性性质

看到“期望”，我们聪明的脑袋瓜就应该想到期望的线性性质，呐！一个复杂随机变量的期望，等于它分解成的简单随机变量的期望之和。

虚树的权值 $w(S)$ 是不同颜色的数量。我们可以把它分解成对每一种颜色的贡献。设 $I_c$ 是一个指示器随机变量，当颜色 $c$ 出现在 $S$ 的虚树中时， $I_c=1$ ，否则 $I_c=0$ 。那么 $w(S) = \sum_{c} I_c$ 。

根据期望的线性性质：

E[w(S)] = E[\sum_c I_c] = \sum_c E[I_c]

而指示器变量的期望就是它取 1 的概率，所以：

E[w(S)] = \sum_c P(\text{颜色 } c \text{ 在虚树中})

第二步：颜色 c 出现的条件

现在问题转化成：对于一个固定的颜色 $c$ 和固定的集合大小 $i$ ，随机选一个点集 $S$ ，它的虚树包含颜色 $c$ 的概率是多少？

直接计算一个颜色出现的概率有点麻烦，我们不如反过来想，计算它不出现的概率，也就是所谓的正难则反，喵~

颜色 $c$ 不出现在虚树中，意味着虚树里所有的边颜色都不是 $c$ 。这等价于，我们选择的点集 $S$ 中的所有点，都位于原树移除所有颜色为 $c$ 的边后，形成的某一个连通块内。

举个例子，如果把所有颜色为 $c$ 的边都剪断，树可能会分裂成好几个连通块。假设分成了 $k_c$ 个块，它们的大小分别是 $s_{c,1}, s_{c,2}, \dots, s_{c,k_c}$ 。如果点集 $S$ 的所有点都来自同一个块 $T_{c,j}$ ，那么连接它们的任何路径都不会经过颜色为 $c$ 的边，虚树里自然就没有颜色 $c$ 啦。

所以，对于大小为 $i$ 的点集 $S$ ：

从块 $T_{c,j}$ 中选出 $i$ 个点的方法数是 $\binom{s_{c,j}}{i}$ 。
使得颜色 $c$ 不出现的总方法数是 $\sum_{j=1}^{k_c} \binom{s_{c,j}}{i}$ 。
随机选大小为 $i$ 的点集的总方法数是 $\binom{n}{i}$ 。

于是，颜色 $c$ 不出现的概率是：

P(\text{颜色 } c \text{ 不在虚树中}) = \frac{\sum_{j=1}^{k_c} \binom{s_{c,j}}{i}}{\binom{n}{i}}

那么，颜色 $c$ 出现的概率就是 $1$ 减去上面这个值。

第三步：汇总所有颜色

把所有颜色加起来，我们得到大小为 $i$ 的点集权值的期望总和 $W_i = \sum_{|S|=i} w(S)$ ：

W_i = \sum_c \left( \binom{n}{i} - \sum_{j=1}^{k_c} \binom{s_{c,j}}{i} \right) = (\text{总颜色数}) \binom{n}{i} - \sum_c \sum_j \binom{s_{c,j}}{i}

我们要求的期望 $E_i = \frac{W_i}{\binom{n}{i}}$ 。

这个式子对每个 $i$ 单独计算还是太慢了。我们需要一个能一次性算出所有 $i$ 的答案的方法。这通常指向了生成函数和多项式，喵！

第四步：神奇的树形 DP 和容斥

核心问题变成了如何高效地计算 $\sum_c \sum_j \binom{s_{c,j}}{i}$ 这一项。如果我们能对每个 $k$ 求出 $\alpha_k$ （即在所有颜色 $c$ 的划分中，大小为 $k$ 的连通块一共出现了多少次），那么 $\sum_c \sum_j \binom{s_{c,j}}{i} = \sum_{k=i}^n \alpha_k \binom{k}{i}$ 。

但是直接求 $\alpha_k$ 依然困难。这里有一个非常巧妙的思路，结合了树形 DP 和容斥思想。

我们不直接计算最终的连通块大小，而是计算每次分割操作对 $\sum \binom{\text{size}}{i}$ 的贡献。

初始时，对于每一种颜色，我们都只有一个大小为 $n$ 的连通块。
我们对原树进行一次 DFS。当我们经过一条边 $(u, v)$ （ $u$ 是 $v$ 的父节点），颜色为 $c$ ，这条边会把一个原本的连通块（对于颜色 $c$ 而言）分割成两部分。
假设分割前这个块的大小是 $S$ ，分割后变成了大小为 $S_1$ 和 $S_2$ 的两个块（这里 $S_1$ 是 $v$ 的子树大小 sz[v], $S_2 = S - sz[v]$ ）。
这个分割操作对 $\sum \binom{\text{size}}{i}$ 的贡献改变是 $\left(\binom{S_1}{i} + \binom{S_2}{i}\right) - \binom{S}{i}$ 。

把所有分割的贡献累加起来，就能得到 $\sum_c \sum_j \binom{s_{c,j}}{i}$ 的总和。设 $U_i = \sum_c \sum_j \binom{s_{c,j}}{i}$ 。 $U_i = (\text{总颜色数}) \binom{n}{i} + \sum_{\text{所有分割 } S \to S_1, S_2} \left(\binom{S_1}{i} + \binom{S_2}{i} - \binom{S}{i}\right)$

我们可以定义一个系数数组 $A$ ，其中 $A_k$ 表示 $\binom{k}{i}$ 的系数。

初始时，对于每种颜色，我们有一个大小为 $n$ 的块，所以 $A_n$ 加上总颜色数。
对于每次分割 $S \to S_1, S_2$ ，我们从 $A$ 中减去一个 $S$ 的贡献，加上 $S_1$ 和 $S_2$ 的贡献。即 $A_S \to A_S-1$ , $A_{S_1} \to A_{S_1}+1$ , $A_{S_2} \to A_{S_2}+1$ 。

这样，我们就可以通过一次 DFS 遍历所有边，模拟这个分割过程，从而计算出最终的系数数组 $A$ 。 $U_i = \sum_{k=0}^n A_k \binom{k}{i}$ 。

而我们要求的 $W_i = (\text{总颜色数})\binom{n}{i} - U_i = (\text{总颜色数})\binom{n}{i} - \sum_{k=0}^n A_k \binom{k}{i}$ 。把 $A_k$ 的定义代入，可以发现一个美妙的抵消： $W_i = -\sum_{\text{所有分割 } S \to S_1, S_2} \left(\binom{S_1}{i} + \binom{S_2}{i} - \binom{S}{i}\right)$ 令 $B_k$ 为 $\binom{k}{i}$ 的系数，那么对于每次分割 $S \to S_1, S_2$ ，我们有 $B_S \to B_S+1, B_{S_1} \to B_{S_1}-1, B_{S_2} \to B_{S_2}-1$ 。我们可以通过一次 DFS 来计算出最终的系数数组 $B$ 。

第五步：多项式乘法 (NTT)

现在我们有了 $W_i = \sum_{k=0}^n B_k \binom{k}{i}$ 。这是一个典型的可以被多项式乘法优化的形式。展开组合数：

W_i = \sum_{k=i}^n B_k \frac{k!}{i!(k-i)!}

整理一下：

\frac{W_i}{i!} = \sum_{k=i}^n (B_k \cdot k!) \cdot \frac{1}{(k-i)!}

这是一个卷积的形式！令 $P_k = B_k \cdot k!$ ， $Q_k = \frac{1}{k!}$ 。我们想求的 $C_i = \frac{W_i}{i!}$ 可以通过计算两个多项式的乘积得到。具体地，令多项式 $P(x) = \sum_{k=0}^n P_{n-k} x^k$ 和 $Q(x) = \sum_{k=0}^n Q_k x^k$ 。它们的乘积 $R(x) = P(x)Q(x)$ 中 $x^{n-i}$ 项的系数，恰好就是 $C_i$ ！

所以，完整流程是：

DFS 一遍，计算每个节点的子树大小 sz。
再 DFS 一遍，模拟分割过程，计算出系数数组 $B_k$ 。
构造多项式 $P(x)$ 和 $Q(x)$ 。
使用 NTT（快速数论变换）计算它们的乘积 $R(x)$ 。
从 $R(x)$ 的系数中提取出 $W_i$ 。
最后，计算期望 $E_i = W_i / \binom{n}{i}$ 。

好耶！这样我们就把一个复杂的期望问题转化成了我们熟悉的树形 DP 和多项式板子题啦，喵~

代码实现

cpp

#include <iostream>
#include <vector>
#include <numeric>
#include <algorithm>

using namespace std;

const int MOD = 998244353;
const int G = 3;

// 快速幂，喵~
long long power(long long base, long long exp) {
    long long res = 1;
    base %= MOD;
    while (exp > 0) {
        if (exp % 2 == 1) res = (res * base) % MOD;
        base = (base * base) % MOD;
        exp /= 2;
    }
    return res;
}

// 模逆元，喵~
long long modInverse(long long n) {
    return power(n, MOD - 2);
}

// NTT 核心实现，就像我爪子一样快！
void ntt(vector<long long>& a, bool invert) {
    int n = a.size();
    for (int i = 1, j = 0; i < n; i++) {
        int bit = n >> 1;
        for (; j & bit; bit >>= 1)
            j ^= bit;
        j ^= bit;
        if (i < j)
            swap(a[i], a[j]);
    }

    for (int len = 2; len <= n; len <<= 1) {
        long long wlen = power(G, (MOD - 1) / len);
        if (invert)
            wlen = modInverse(wlen);
        for (int i = 0; i < n; i += len) {
            long long w = 1;
            for (int j = 0; j < len / 2; j++) {
                long long u = a[i + j], v = (a[i + j + len / 2] * w) % MOD;
                a[i + j] = (u + v) % MOD;
                a[i + j + len / 2] = (u - v + MOD) % MOD;
                w = (w * wlen) % MOD;
            }
        }
    }

    if (invert) {
        long long n_inv = modInverse(n);
        for (long long& x : a)
            x = (x * n_inv) % MOD;
    }
}

const int MAXN = 200005;
vector<pair<int, int>> adj[MAXN];
int subtree_size[MAXN];
long long color_comp_size[MAXN];
long long poly_coeffs[MAXN * 2]; // B_k 数组

// 第一次 DFS：计算子树大小
void dfs_size(int u, int p) {
    subtree_size[u] = 1;
    for (auto& edge : adj[u]) {
        int v = edge.first;
        if (v == p) continue;
        dfs_size(v, u);
        subtree_size[u] += subtree_size[v];
    }
}

// 第二次 DFS：计算 B_k 系数
void dfs_build_coeffs(int u, int p) {
    for (auto& edge : adj[u]) {
        int v = edge.first;
        int color = edge.second;
        if (v == p) continue;

        long long S = color_comp_size[color];
        long long S1 = subtree_size[v];
        long long S2 = S - S1;

        poly_coeffs[S] = (poly_coeffs[S] + 1) % MOD;
        poly_coeffs[S1] = (poly_coeffs[S1] - 1 + MOD) % MOD;
        poly_coeffs[S2] = (poly_coeffs[S2] - 1 + MOD) % MOD;

        long long old_S_for_color = color_comp_size[color];
        color_comp_size[color] = S1;
        dfs_build_coeffs(v, u);
        color_comp_size[color] = old_S_for_color; // 回溯
    }
}


int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);

    int n;
    cin >> n;

    for (int i = 1; i < n; ++i) {
        int u, v, c;
        cin >> u >> v >> c;
        adj[u].push_back({v, c});
        adj[v].push_back({u, c});
    }

    // 预处理
    dfs_size(1, 0);
    for (int i = 1; i <= n; ++i) {
        color_comp_size[i] = n;
    }

    // 计算多项式系数
    dfs_build_coeffs(1, 0);

    vector<long long> fact(n + 1);
    vector<long long> invFact(n + 1);
    fact[0] = 1;
    invFact[0] = 1;
    for (int i = 1; i <= n; ++i) {
        fact[i] = (fact[i - 1] * i) % MOD;
        invFact[i] = modInverse(fact[i]);
    }

    // 准备 NTT
    int len = 1;
    while (len <= 2 * n) len <<= 1;

    vector<long long> P(len, 0), Q(len, 0);
    for (int k = 0; k <= n; ++k) {
        P[k] = (poly_coeffs[n - k] * fact[n - k]) % MOD;
    }
    for (int k = 0; k <= n; ++k) {
        Q[k] = invFact[k];
    }

    ntt(P, false);
    ntt(Q, false);
    for (int i = 0; i < len; ++i) {
        P[i] = (P[i] * Q[i]) % MOD;
    }
    ntt(P, true);

    // 计算并输出最终答案
    for (int i = 1; i <= n; ++i) {
        long long Wi_times_inv_i_fact = P[n - i];
        long long Wi = (Wi_times_inv_i_fact * fact[i]) % MOD;
        
        long long C_n_i = (fact[n] * invFact[i]) % MOD;
        C_n_i = (C_n_i * invFact[n - i]) % MOD;
        
        long long Ei = (Wi * modInverse(C_n_i)) % MOD;
        cout << Ei << (i == n ? "" : " ");
    }
    cout << endl;

    return 0;
}

复杂度分析

时间复杂度: $O(N \log N)$
- 两次 DFS 遍历树都是 $O(N)$ 的。
- 构造多项式是 $O(N)$ 。
- 核心计算是 NTT，其长度为 $O(N)$ ，所以时间复杂度是 $O(N \log N)$ 。
- 预处理阶乘和逆元是 $O(N)$ 。
- 因此，总的时间复杂度由 NTT 决定，为 $O(N \log N)$ ，喵~
空间复杂度: $O(N)$
- 邻接表、子树大小等数组需要 $O(N)$ 的空间。
- NTT 需要的数组长度是 $O(N)$ 的。
- 所以总的空间复杂度是 $O(N)$ ，非常优秀的说！

知识点总结

这道题是多种算法思想的完美结合，做完之后感觉自己的小脑袋瓜都变聪明了呢，喵~

期望的线性性质: 解决复杂期望问题的关键钥匙，将问题分解为对每个颜色贡献的求和。
贡献法/容斥原理: 通过分析每一次操作（分割连通块）对最终结果的贡献，将一个静态的计数问题转化为一个动态的、可累加的过程。
树形 DP: 两次 DFS 都利用了树的递归结构。第一次计算子树大小，第二次巧妙地利用递归和回溯来维护和计算每个颜色的连通块分割情况。
组合数学: 解题过程大量使用了组合数 $\binom{n}{k}$ 及其性质。
多项式与 NTT: 最终的求和式 $\sum B_k \binom{k}{i}$ 是一个卷积形式，使用 NTT 可以从 $O(N^2)$ 优化到 $O(N \log N)$ ，是解决这类计数问题的强大武器。

希望这篇题解能帮助你理解这道有趣的题目！如果还有不明白的地方，随时可以再来问我哦，喵~ >w<

S 老师的虚树 - 题解 ​

标签与难度 ​

题目大意喵~ ​

解题思路分析 ​

第一步：期望的线性性质 ​

第二步：颜色 c 出现的条件 ​

第三步：汇总所有颜色 ​

第四步：神奇的树形 DP 和容斥 ​

第五步：多项式乘法 (NTT) ​

代码实现 ​

复杂度分析 ​

知识点总结 ​