[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();}
}