P14254 树上组合计数（分割问题）题解
原题链接
一、题目分析
这部分是解题的核心，通过分析条件得出简化问题的关键结论。
-
计数问题先尝试找一下性质：
-
注意到节点的选择只能越来越深
\[d_i>=d_1 \] -
最关键性质：
\[max L_i =d_i ,min R_i ==high_i \]\[L_i =d_i<=d_1,\therefore d_i=d_1 \]
-
-
所以所选点的深度相等
- 定义“高度”hig[x]：结点x到其所在子树中叶子结点的最大深度（即子树内的最大深度）。
- 核心约束：设选中结点均在深度为x的层，需满足“所有选中结点的高度最小值等于第一个结点的高度”（min Rᵢ = R₁，其中Rᵢ为结点i的高度）。
三、计数思路
基于上述性质，计数过程按“逐层处理”展开，每层内再按高度分组计算贡献。
1. 预处理步骤
- 步骤1：计算深度与高度：通过DFS遍历树，记录每个结点的深度dep[x]和高度hig[x]。
- 步骤2：分层与分组：
- 按深度将结点分组，得到flr[i]（存储深度为i的所有结点的高度）。
- 对每层的flr[i]按高度排序，并进一步按相同高度分组，记录每组的结点数量（nt）和该组在层中的位置（bef，用于计算后续结点数）。
2. 每层贡献计算
对每个深度为i的层，若该层结点数<k+1，则直接跳过（无法满足选k个结点的基本条件）；否则按以下公式计算该层的贡献：
（1）核心公式
对于高度相同的一组结点（数量为t，后续结点数为s，s=该层总结点数 - 该组位置bef）：
- 若满足t≥2（需至少2个结点，1个作为序列第一个元素，1个用于限制高度最小值）且t+s≥k+1（总可选结点数足够），则该组贡献为：
[ \text{贡献} = t \times \left( A(t+s-1, k-1) - A(s, k-1) \right) ]
- 特殊情况：当s=k-1时，A(s, k-1)需强制设为0（此时后续结点数恰好为k-1，无法满足高度约束）。
（2）排列数定义
- 排列数A(a, b)表示从a个元素中选b个元素进行有序排列的方案数。
- 计算公式：若b<0或a<b，则A(a,b)=0；否则A(a,b)=fac[a] × inf[a-b] mod 998244353。
- 其中fac为阶乘数组，inf为逆阶乘数组（通过费马小定理预处理，inf[n] = fac[n]^(M-2) mod M，M=998244353）。
四、代码实现
1.完整代码
#include <bits/stdc++.h>
using namespace std;
#define int long long
const int N = 1e6 + 5, M = 998244353;
struct Edge
{int to, next;
} e[N];
int head[N], idx = 0;
void addedge(int u, int v)
{e[idx].to = v;e[idx].next = head[u];head[u] = idx++;
}
int phi(int a, int b)
{int res = 1;while (b){if (b & (int)1)res = (res * a) % M;b >>= (int)1;a = (a * a) % M;}return res;
}
int n, k;
int dep[N], hig[N];
vector<int> flr[N];
int fac[N], inf[N];
void init()
{fac[1] = fac[0] = 1;for (int i = 1; i <= n; ++i)fac[i] = (fac[i - 1] * i) % M;inf[n] = phi(fac[n], M - 2);for (int i = n - 1; i >= 0; --i)inf[i] = (inf[i + 1] * (i + 1)) % M;
}
int maxdep = 0;
void dfs(int x)
{hig[x] = dep[x];maxdep = max(maxdep, dep[x]);for (int i = head[x]; ~i; i = e[i].next){int v = e[i].to;dep[v] = dep[x] + 1;dfs(v);hig[x] = max(hig[x], hig[v]);}
}
int A(int a, int b)
{if (b < 0 || a < b)return 0;return (fac[a] * inf[a - b]) % M;
}
struct Seg
{int nt = 0, bef = 0;
};
signed main()
{memset(head, -1, sizeof head);cin >> n >> k;for (int i = 2; i <= n; ++i){int t;cin >> t;addedge(t, i);}dep[1] = 1;dfs(1);init();for (int i = 1; i <= n; ++i)flr[dep[i]].push_back(hig[i]);int ans = 0;for (int i = 1; i <= maxdep; ++i){if (flr[i].size() < k + 1)continue;sort(flr[i].begin(), flr[i].end());vector<Seg> g;int cnt = 1;for (int j = 0; j < flr[i].size(); ++j){if (j < flr[i].size() && flr[i][j] == flr[i][j + 1])++cnt;elseg.push_back({cnt, j + 1}), cnt = 1;}for (Seg j : g){int t = j.nt, s = flr[i].size() - j.bef;if (t >= 2 && t + s >= k + 1){int t1 = A(t + s - 1, k - 1), t2 = A(s, k - 1);if (s == k - 1)t2 = 0;int tp = ((t1 - t2) % M + M) % M;ans = (ans + (t * tp) % M) % M;}}}cout << ans << endl;return 0;
}
2. 易错点提示
排列数函数记得特判
3. 代码说明
- 树的存储：使用邻接表（Edge结构体+head数组），适配1e6级别的结点规模。
- 预处理优化：阶乘与逆阶乘通过费马小定理预处理，确保排列数计算O(1)。
- 边界处理：
- 排列数计算时判断a<b或b<0的情况，直接返回0。
- 减法取模时加上M再取模，避免出现负数。