E - Count 123 畳み込み解法

著者コメント

想定に対して高度なライブラリを使用しますが、実際の考察の負荷は低いと思います。

※ 文章中の \((\ \ )\) は組合せ（二項係数）を表し、

\[ \binom{n}{k} = \frac{n!}{k!(n-k)!} \]

を意味します。

立式

まず 2 を先に並べる。

このとき、2 の前後にある隙間は全部で X_2+1 個ある。

_ 2 _ 2 _ ... _ 2 _

1 と 3 は隣り合えないので、同じ隙間に 1 と 3 の両方を入れることはできない。

そこで、使う隙間の個数を S とする。

さらに、その S 個の隙間のうち、1 を入れる隙間の個数を P とする。

すると、3 を入れる隙間の個数は S-P 個である。

固定した S, P に対する寄与を考える。

まず、X_2+1 個の隙間から使う S 個を選ぶ方法は

\[ \binom{X_2+1}{S} \]

次に、その S 個のうち P 個を 1 用にする方法は

\[ \binom{S}{P} \]

通りである。

残りの S-P 個の隙間は 3 用になる。

X_1 個の 1 を P 個の非空な塊に分ける方法は

\[ \binom{X_1-1}{P-1} \]

通りである。

同様に、X_3 個の 3 を S-P 個の非空な塊に分ける方法は

\[ \binom{X_3-1}{S-P-1} \]

通りである。

よって、固定した S, P に対する寄与は

\[ \binom{X_2+1}{S} \binom{S}{P} \binom{X_1-1}{P-1} \binom{X_3-1}{S-P-1} \]

である。

したがって、答えは

\[ \sum_S \sum_P \binom{X_2+1}{S} \binom{S}{P} \binom{X_1-1}{P-1} \binom{X_3-1}{S-P-1} \]

である。

---

畳み込みの形にする

Q = S-P とおく。

このとき、

\[ \binom{S}{P}= \frac{S!}{P!Q!} \]

である。

したがって、固定した S, P に対する寄与は

\[ \binom{X_2+1}{S} S! \left( \frac{\binom{X_1-1}{P-1}}{P!} \right) \left( \frac{\binom{X_3-1}{Q-1}}{Q!} \right) \]

と変形できる。

ここで

\[ A_P = \frac{\binom{X_1-1}{P-1}}{P!} \]

\[ B_Q = \frac{\binom{X_3-1}{Q-1}}{Q!} \]

とおく。

S を固定したとき、P+Q=S となるすべての組を足せばよいので、

\[ C_S = \sum_{P+Q=S} A_PB_Q \]

と書ける。

これは A と B の畳み込みそのものである。

よって C = convolution(A, B) とすればよい。

最終的な答えは

\[ \sum_S \binom{X_2+1}{S} S! C_S \]

である。

---

計算量

階乗と階乗逆元を前計算しておけば、二項係数は O(1) で求められる。

畳み込みを NTT で行うので、全体の計算量は

\[ O((X_1+X_3)\log(X_1+X_3)) \]

である。

---

実装例 C++

コードは解説用に改めて書いたものです。

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

using mint = modint998244353;

int main() {
    ios::sync_with_stdio(false);
    cin.tie(nullptr);

    int X1, X2, X3;
    cin >> X1 >> X2 >> X3;

    int N = max({X1, X2 + 1, X3}) + 5;

    vector<mint> fact(N), ifact(N);

    fact[0] = 1;

    for (int i = 1; i < N; i++) {
        fact[i] = fact[i - 1] * i;
    }

    ifact[N - 1] = fact[N - 1].inv();

    for (int i = N - 1; i >= 1; i--) {
        ifact[i - 1] = ifact[i] * i;
    }

    auto C = [&](int n, int k) -> mint {
        if (k < 0 || k > n) return 0;
        return fact[n] * ifact[k] * ifact[n - k];
    };

    vector<mint> A(X1 + 1), B(X3 + 1);

    for (int p = 1; p <= X1; p++) {
        A[p] = C(X1 - 1, p - 1) * ifact[p];
    }

    for (int q = 1; q <= X3; q++) {
        B[q] = C(X3 - 1, q - 1) * ifact[q];
    }

    vector<mint> conv = convolution(A, B);

    mint ans = 0;

    for (int s = 0; s <= X2 + 1 && s < (int)conv.size(); s++) {
        ans += C(X2 + 1, s) * fact[s] * conv[s];
    }

    cout << ans.val() << '\n';
}