数树（2021 CoE-II E）

题目链接

前置知识

（1）生成函数（解决组合或排列问题的一种有力工具）：

• 浅谈生成函数之OGF
• Ordinary Generating Function

（2）拉格朗日反演（根据复合逆得到函数的某项系数）：

• 拉格朗日反演
• Lagrange Inversion Formula

正文

浏览一遍题目，可知等概率生成一棵符合题目约束的树，其 v(\mathcal T) 的期望值等于所有树的权值 v 之和除以不同构树的总数 c。由于 a，b 均为整数，可知 v 必定是整数，而 c 显然是整数，当 n 较小时，可以暴力生成所有可能的树，分别统计权值和得到答案。当 n 较大时，需要寻找高效的方法（附注：测试点 1 是给暴力生成树的解题者设置，测试点 2 是给暴力生成树并加以适当剪枝的解题者设置，测试点 3 和 4 是为潜在的 \mathcal O(n^2)动态规划设置；测试点 5 和 6 留给愿意写多项式算法的解题者，时间 \mathcal O(n\log n)，要注意常数）。

首先给出结论：

\large \mathbb{E}(v(\mathcal T)) = \frac{[x^{n-1}]\left(\left(1+x^{k_1}+x^{k_2}\right)^{n-1}\left(ax^{k_1}+bx^{k_2}\right)\right)}{\dfrac{1}{n}[x^{n-1}]\left(1+x^{k_1}+x^{k_2}\right)^n}

先确定分母。令符合题意要求的由 n 个点组成的不同构的树的数量所构成序列的普通生成函数为 \mathbf F(x)，由题意，只有 1 个结点的任意 k_1-k_2~\mathrm{Tree} 的数量为 1，因此 \mathbf F(x) 的一次幂 x 的系数为 1，根结点有 k_1 棵子树（或者 k_2 棵子树），这 k_1 棵子树（或者 k_2 棵子树）总的结点数为 n - 1，且子树也是 k_1-k_2~\mathrm{Tree}，结合判定同构的规则， 可以得到：

\mathbf F(x)=x\left(1+\mathbf F^{k_1}(x)+\mathbf F^{k_2}(x)\right)

观察可知 \mathbf F(x) 的零次项系数为 0，一次项系数为 1，且：

x=\dfrac{x\left(1+\mathbf F^{k_1}(x)+\mathbf F^{k_2}(x)\right)}{1+\mathbf F^{k_1}(x)+\mathbf F^{k_2}(x)}=\dfrac{\mathbf F(x)}{1+\mathbf F^{k_1}(x)+\mathbf F^{k_2}(x)}

因此其复合逆为：

\mathbf{F}^{\langle-1\rangle}(x)=\dfrac{x}{1+x^{k_1}+x^{k_2}}

根据拉格朗日反演，有：

[x^n]\mathbf F(x)=\dfrac{1}{n}[x^{n-1}]\left(1+x^{k_1}+x^{k_2}\right)^n

正是分母。

现在要考虑求所有树的权值和，一个常见的思路是使用多元生成函数，然后求导。

令

\mathbf F(x,\ z)=\displaystyle \sum f_{n,m}x^nz^m

其中 f_{n,m} 表示 n 个结点的 k_1-k_2~\mathrm{Tree} 权值和为 m 的方案数，由于根结点要么有 k_1 棵子树，要么有 k_2 棵子树，要么没有子树。如果有 k_1 棵子树，加上根结点后，又构成了一个 k_1 叉结点，每种方案增加了 a 分。如果有 k_2 棵子树，加上根结点后，又构成了一个 k_2 叉结点，每种方案增加了 b 分。如果树中只包含根结点，则由于 1 个结点的 k_1-k_2~\mathrm{Tree} 权值和为 0 的方案数为 1，故有：

\mathbf F(x,\ z)=x\left( 1 + z^{a}\mathbf F^{k_1}(x,z) + z^b\mathbf F^{k_2}(x,z) \right)

类似地，容易发现 x 这一维的复合逆是:

\mathbf F_x^{\langle -1\rangle}(x,\ z)=\dfrac{x}{1+z^{a}x^{k_1}+z^bx^{k_2}}

显然多元生成函数也可以进行拉格朗日反演：

[x^n]\mathbf F(x,\ z)=\frac{1}{n}[x^{n-1}]\left( 1+z^ax^{k_1}+z^{b}x^{k_2} \right)^n

令：

\zeta (x,\ z)=\left(1+z^ax^{k_1}+z^bx^{k_2}\right)^n

现在要求权值和，对 z 这一维求导：

\frac{\partial\zeta(x,\ z)}{\partial z}=n\left( 1+z^ax^{k_1}+z^bx^{k_2} \right)^{n-1}\left(az^{a-1}x^{k_1}+ bz^{b-1}x^{k_2} \right)

显然权值和就是：

[x^n]\frac{\partial \mathbf F(x,\ z)}{\partial z}\biggr\vert_{z=1}=\frac{1}{n}[x^{n-1}]\frac{\partial\zeta(x,\ z)}{\partial z}\biggr\vert_{z=1}=[x^{n-1}]\left ( \left( 1 +x^{k_1}+x^{k_2} \right)^{n-1}\left( ax^{k_1}+bx^{k_2} \right) \right)

可以 \mathcal O(n) 枚举 x^{k_1} 选了几次，用多项式系数计算答案即可。由于需要求同余方程的解，可以将求逆元与解题代码有机结合。

出题人给出的标程：

#include <bits/stdc++.h>
using namespace std;

const int N = 1e7 + 5, mod = 998244353;

inline int inc(int a, int b) { return (a += b - mod) + (a >> 31 & mod); }
inline int dec(int a, int b) { return (a -= b) + (a >> 31 & mod); }
inline int mul(int a, int b) { return 1ll * a * b % mod; }
inline void Inc(int &a, int b) { return void(a = inc(a, b)); }
inline void Dec(int &a, int b) { return void(a = dec(a, b)); }
inline void Mul(int &a, int b) { return void(a = mul(a, b)); }
inline int Pow(int a, int b) {
    int c = 1;
    for (; b; Mul(a, a), b >>= 1)
        if (1 & b)
            Mul(c, a);
    return c;
}
inline int Inv(int x) { return Pow(x, mod - 2); }

int fac[N], inv[N], n, k1, k2, a, b, tot, Sum;

inline int Polycoef(int n, int m, int k) { return mul(fac[n + m + k], mul(mul(inv[n], inv[m]), inv[k])); }

int main(void) {
    scanf("%d%d%d%d%d", &k1, &k2, &n, &a, &b);
    fac[0] = 1;
    for (int i = 1; i <= n; i++) fac[i] = mul(fac[i - 1], i);
    inv[1] = 1;
    for (int i = 2; i <= n; i++) inv[i] = mul(mod - mod / i, inv[mod % i]);
    inv[0] = 1;
    for (int i = 1; i <= n; i++) Mul(inv[i], inv[i - 1]);
    for (int i = 0; i <= n; i++)
        if (i * k1 <= n - 1) {
            int res = n - 1 - i * k1;
            if (res % k2)
                continue;
            int j = res / k2;
            if (j > n)
                continue;
            Inc(tot, Polycoef(i, j, n - i - j));
        }
    for (int i = 0; i <= n - 1; i++)
        if (i * k1 <= n - 1 - k1) {
            int res = n - 1 - k1 - i * k1;
            if (res % k2)
                continue;
            int j = res / k2;
            if (j > n - 1)
                continue;
            Inc(Sum, mul(Polycoef(i, j, n - 1 - i - j), a));
        }
    for (int i = 0; i <= n - 1; i++)
        if (i * k1 <= n - 1 - k2) {
            int res = n - 1 - k2 - i * k1;
            if (res % k2)
                continue;
            int j = res / k2;
            if (j > n - 1)
                continue;
            Inc(Sum, mul(Polycoef(i, j, n - 1 - i - j), b));
        }
    Mul(tot, Inv(n));
    return printf("%d\n", mul(Sum, Inv(tot))) * 0;
}