[2023杭电多校5 1005] Snake (生成函数)

题意

有 $n$ 个标号为 $1,2,\cdots ,n$ 的球，放到 $m$ 个无标号盒子 (盒内顺序有标号)，且每个盒子球数不超过 $k$，求方案数对 $998244353$ 取模。

$1\le m,k\le n\le 10^6$

分析：

考虑每个盒子内球的生成函数 $\sum _{i=1}^k{x}^i$，那么 $m$ 个盒子的生成函数就为 ${\left(\sum _{i=1}^k{x}^i\right)}^m$，那么方案数就为第 $n$ 项系数

由于球带标号，所以需要对答案全排列，也就是乘 $n!$，又由于盒子不带标号，所以要对答案除 $m!$，那么答案为

$$\frac{n!}{m!}\times [x^n]{\left(\sum _{i=1}^k{x}^i\right)}^m$$

$10^6$ 用多项式快速幂会超时，考虑

$${\left(\sum _{i=1}^k{x}^i\right)}^m=x^m{\left(\sum _{i=0}^{k-1}{x}^i\right)}^m=x^m\frac{(1-x^k{)}^m}{(1-x{)}^m}$$

转为求 $[x^{n-m}]\frac{(1-x^k{)}^m}{(1-x{)}^m}$ 其中

$$\begin{gathered} (1-x^k{)}^m=\sum _{i=0}^m(\genfrac{}{}{0ex}{}{m}{i})\times (-1{)}^i\times x^{i\times k} \\ \frac{1}{(1-x{)}^m}=\sum _{i=0}^{\infty }\left(\genfrac{}{}{0ex}{}{m-1+i}{m-1}\right)\times x^i \end{gathered}$$

于是枚举第一个式子的 $i$，那么只需要求第二个式子的 $n-m-i\times k$ 项系数即可。预处理组合数即可。

代码：

#include <bits/stdc++.h>
using namespace std;
using i64 = long long;
constexpr int mod = 998244353;
template<class T>
T power(T a, int b) {
    T res = 1;
    for (; b; b /= 2, a *= a) {
        if (b % 2) {
            res *= a;
        }
    }
    return res;
}
template<int mod>
struct ModInt {
    int x;
    ModInt() : x(0) {}
    ModInt(i64 y) : x(y >= 0 ? y % mod : (mod - (-y) % mod) % mod) {}
    ModInt &operator+=(const ModInt &p) {
        if ((x += p.x) >= mod) x -= mod;
        return *this;
    }
    ModInt &operator-=(const ModInt &p) {
        if ((x += mod - p.x) >= mod) x -= mod;
        return *this;
    }
    ModInt &operator*=(const ModInt &p) {
        x = (int)(1LL * x * p.x % mod);
        return *this;
    }
    ModInt &operator/=(const ModInt &p) {
        *this *= p.inv();
        return *this;
    }
    ModInt operator-() const {
        return ModInt(-x);
    }
    ModInt operator+(const ModInt &p) const {
        return ModInt(*this) += p;
    }
    ModInt operator-(const ModInt &p) const {
        return ModInt(*this) -= p;
    }
    ModInt operator*(const ModInt &p) const {
        return ModInt(*this) *= p;
    }
    ModInt operator/(const ModInt &p) const {
        return ModInt(*this) /= p;
    }
    bool operator==(const ModInt &p) const {
        return x == p.x;
    }
    bool operator!=(const ModInt &p) const {
        return x != p.x;
    }
    ModInt inv() const {
        int a = x, b = mod, u = 1, v = 0, t;
        while (b > 0) {
            t = a / b;
            swap(a -= t * b, b);
            swap(u -= t * v, v);
        }
        return ModInt(u);
    }
    ModInt pow(i64 n) const {
        ModInt res(1), mul(x);
        while (n > 0) {
            if (n & 1) res *= mul;
            mul *= mul;
            n >>= 1;
        }
        return res;
    }
    friend ostream &operator<<(ostream &os, const ModInt &p) {
        return os << p.x;
    }
    friend istream &operator>>(istream &is, ModInt &a) {
        i64 t;
        is >> t;
        a = ModInt<mod>(t);
        return (is);
    }
    int val() const {
        return x;
    }
    static constexpr int val_mod() {
        return mod;
    }
};
using Z = ModInt<mod>;
vector<Z> fact, infact;
void init(int n) {
    fact.resize(n + 1), infact.resize(n + 1);
    fact[0] = infact[0] = 1;
    for (int i = 1; i <= n; i ++) {
        fact[i] = fact[i - 1] * i;
    }
    infact[n] = fact[n].inv();
    for (int i = n; i; i --) {
        infact[i - 1] = infact[i] * i;
    }
}
Z C(int n, int m) {
    if (n < 0 || m < 0 || n < m) return Z(0);
    return fact[n] * infact[n - m] * infact[m];
}
void solve() {
    int n, m, k;
    cin >> n >> m >> k;
    Z ans;
    for (int i = 0; i <= m; i ++) {
        Z f = i & 1 ? Z(-1) : Z(1);
        ans += f * C(m, i) * C(n - k * i - 1, m - 1);
    }
    cout << ans * fact[n] / fact[m] << "\n";
}
signed main() {
    init(1e6);
    cin.tie(0) -> sync_with_stdio(0);
    int T;
    cin >> T;
    while (T --) {
        solve();
    }
}