P14254 分割（树上计数问题） 题解

P14254 树上组合计数（分割问题）题解

原题链接

一、题目分析

这部分是解题的核心，通过分析条件得出简化问题的关键结论。

1. 计数问题先尝试找一下性质：

   • 注意到节点的选择只能越来越深

     \[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 \]

2. 所以所选点的深度相等

   • 定义“高度”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;
            else
                g.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再取模，避免出现负数。