「FPS 24 題」全部解く: C「数列」

解きます

解く

「$M$以下の非負整数」に対応する関数は
$$1+x+\dots+x^M=\frac{1-x^{M+1}}{1-x}$$
です. これが$N$個で総和が$S$なので求めるものは
$$\left[x^S\right]\left(\frac{1-x^{M+1}}{1-x}\right)^N$$
になります.

これをどう求めるかですが, とりあえず
$$\left(\frac{1-x^{M+1}}{1-x}\right)^N = \left(1-x^{M+1}\right)^N \frac{1}{\left(1-x\right)^N}$$
という風に分解してその積で求めることにします.

まず, $\left(1-x^{M+1}\right)^N$の側は二項定理を使えば簡単に求められます.
$$
\left[x^A\right]\left(1-x^{M+1}\right)^N = \begin{cases}
\left(-1\right)^{A\div\left(M+1\right)}\binom{N}{A\div\left(M+1\right)} &\text{if } M+1 \mid A \\
0 &\text{if } M+1 \nmid A
\end{cases}
$$

右側は$\frac{1}{1-x}=1+x+x^2+\dots$を$N$回掛けてるので, $B$次の係数は「$N$個の非負整数(上限無し)の列で和が$B$になるものの個数」という数え上げの問題に言い換えられます.
これを解いて("$B$個のものに$N-1$個仕切りを入れる"のやつが有名でしょうか)答えは
$$
\left[x^B\right]\frac{1}{\left(1-x\right)^N} = \binom{B+N-1}{N-1}
$$
となります.

あとは求めた2つの積を出せばいいですね. $S$次の項だけ分かれば良いのでFFTを持ち出すまでもなく$O(S)$で求められます.

use ac_library::*;
use proconio::*;

type Mint = ModInt998244353;

fn main() {
    input! {
        n: usize,
        m: usize,
        s: usize,
    }
    // 階乗テーブル
    let factorial = {
        let mut r = vec![Mint::raw(1)];
        for i in 1..=400000 {
            r.push(r[i - 1] * i);
        }
        r
    };
    // 二項係数
    let binom = |a: usize, b: usize| factorial[a] / (factorial[b] * factorial[a - b]);
    let mut ans = Mint::raw(0);
    // (1-x^(M+1))^Nの0でない次数だけ列挙
    for a in (0..=s).step_by(m + 1) {
        // 対応する1/(1-x)^N側の次数
        let b = s - a;
        let i = a / (m + 1);
        // (1-x^(M+1))^N側の係数
        let lhs = binom(n, i) * if i % 2 == 0 { 1 } else { -1 };
        // 1/(1-x)^N側の係数
        let rhs = binom(b + n - 1, n - 1);
        ans += lhs * rhs;
    }
    println!("{ans}");
}

提出(27 ms):

終わり

FPSっぽい雰囲気を醸し出してきました.

以上です