Advent Calendar Contest 2020 18 日目: 塗るめた

Advent Calendar Contest 2020 の 18 日目として、ecasdqina さん作の問題の tester を担当しました。tester 解は、writer 解と少しアプローチが異なっていたので紹介したいと思います。また、コンテスタントの提出を見て発見した、形式的べき級数を一切使用しないアプローチについても書いています。参加してくれた皆さんのコードが大変参考になりました、ありがとうございました。

問題概要

日本語なので原文参照。ProblemId = 5633 - yukicoder

解法 (1)

まず、問題を言い換えます。以下の問題の答えは元の問題の答えと等しいです。

を満たすについて、以下の操作を行ってできるボールの列としてありうるのは何通りかを計算し、その合計を報告せよ。列は、ある位置のボールについて色が異なったり、書いてある値が異なるときに区別される。
1. $1, 2, \ldots, M$ の中から相異なる $K$ 個の整数を選ぶ。
2. 区別できない $L$ 個の赤いボールを用意する。それぞれのボールに対し、上で選んだ $K$ 種類の整数のうちいずれか 1 つを書き込む。このとき、 $K$ 種類全ての整数について、その整数が書かれたボールが少なくとも $1$ つ存在しなければならない。
3. 区別できない $N-L$ 個の青いボールを用意する。それぞれのボールに対し、 $1$ 以上 $M$ 以下の整数のうちいずれか 1 つを書き込む。
4. 合計 $N$ 個のボールを 1 列に並べる。

問題文で問われている操作は「長さ $N$ の数列を作る」→「数列の中から、ちょうど $K$ 種類になるようにいくつか要素を選ぶ」という手順ですが、この言い換えでは「値が $K$ 種類存在する数列を作る」→「作った数列に対していくつか要素を挿入して長さ $N$ にする」という手順をイメージしています。問題文の操作を逆にしたようなものです。

「 $\mathrm{ways}(k, l) :=$ 長さが $L$ であって、 $K$ 種類中 $k$ 種類の値が含まれているような数列の総数」とおき、これを用いると、解は $\sum_{L=K}^{N} \binom{M}{K} \binom{N}{L} M^{N-L} \mathrm{ways}(K, L)$ となります。上の問題で言うところの 2 番目の手順の通り数を、まるまる $\mathrm{ways}(K, L)$ でおいてしまっている、ということです。

さて、 $\mathrm{ways}(K, L)$ が実際に求められないと解けたことにはなりませんが、これはどのように求められるでしょうか？とりあえず二乗を許容することにすると、以下のような DP が立ちます。

$\mathrm{dp}[0][0] = 1$
$\mathrm{dp}[k][l] = (K-k+1) \times \mathrm{dp}[k-1][l-1] + k \times \mathrm{dp}[k][l-1]$

これを高速化することを考えます。この見た目のまま高速化するのは難しいので、多項式として考えます。以下のように、 $\mathrm{dp}[k]$ に関する多項式を定義します。

$F_{k}(x) = \sum_{l} \mathrm{dp}[k][l] x^{l}$

これを用いて先述の DP を表現すると以下のように変形できます。

$F_{k}(x) = (K-k+1) x F_{k-1}(x) + kx F_{k}(x)$
$(1-kx) F_{k}(x) = (K-k+1) x F_{k-1}(x)$
$F_{k}(x) = \frac{(K-k+1) x}{1 - kx} F_{k-1}(x)$

$F_{k}(x)$ と $F_{k-1}(x)$ との間の漸化式が得られました。さらに変形してみます。

$F_{K}(x) = \frac{1 \times 2 \times \cdots \times K \times x^{K}}{(1 - Kx) \cdots (1-x)} F_{0}(x)$
$F_{K}(x) = \frac{K! x^{K}}{\prod_{k=1}^{K}(1 - kx)}$ $\left( \because F_{0}(x) = 1 \right)$

よって、答えは $\sum_{L=K}^{N} \binom{M}{K} \binom{N}{L} M^{N-L} \times \left[ x^{L} \right] \frac{K! x^{K}}{\prod_{k=1}^{K}(1 - kx)}$ となります。分母にある多項式は分割統治のように FFT (NTT) を適用して多項式を作ってから inv を適用すればよく、これらは十分高速に対処可能です。

ソースコード (1)

ライブラリはもっていなかったので beet さんのものを拝借しました (ありがとうございます)

yukicoder.me

ソースコード本体 (折りたたんでいます)

int main() {
    int N, M, K; cin >> N >> M >> K;

    NTT<2> ntt;
    using mint = NTT<2>::M;
    auto conv = [&](auto as, auto bs) { return ntt.multiply(as, bs); };
    FormalPowerSeries<mint> FPS(conv);
    Enumeration<mint> comb;
    comb.init(200010);
    
    auto go = [&](auto &&self, int l, int r) -> vector<mint> {
        if(r - l == 1) {
            return {mint(1), mint(-r)};
        }
        int m = (l + r) / 2;
        return ntt.multiply(self(self, l, m), self(self, m, r));
    };

    vector<mint> A(K+1);
    A[K] = mint(1);
    for(int i=1; i<=K; i++) A[K] *= mint(i);
    vector<mint> B = FPS.inv(go(go, 0, K), N+1);
    vector<mint> X = ntt.multiply(A, B);

    mint ans(0);
    for(int i=K; i<=N; i++) {
        mint ways = X[i];
        ways *= comb.C(M, K);
        ways *= comb.C(N, i) * mint(M).pow(N-i);
        ans += ways;
    }
    cout << ans << endl;
    return 0;
}

解法 (2)

まず、解法 (1) と同様の言い換えをします (赤いボールと青いボールのお話です)。

さて、赤いボールについては、ちょうど $K$ 種類の整数から構成されなければいけませんでした。これについて少し条件を緩く (？) しつつ、包除原理の方針で考えていくことにします。

本来選ぶべき $K$ 種類の整数のうち、 $i$ 種類を 絶対に選ばない ようにし、あとの $K-i$ 種類についてはどうでもよい (選んでも選ばなくてもよい) ときの場合の数を考えます。これを数えるには「赤いボール $K-i$ 種類・青いボール $M$ 種類の中から好きに選ぶことを $N$ 回行うときの通り数」を求め、具体的に $K$ 種類の中からどの $i$ 種類を選択したかを考慮すればよいため、 $\binom{K}{i} (M+K-i)^{N}$ 通りである、と立式できます。

あとは $i$ の偶奇に応じて足したり引いたりすればよいです。赤いボールに対応させる整数をどう選択するかは $\binom{M}{K}$ 通りあるので、それを掛けるのを忘れないようにしてください。

ソースコード (2)

見ての通り、この方針であれば形式的べき級数を一切使わずに通すことができます。

yukicoder.me

ソースコード本体 (折りたたんでいます)

using mint = ModInt<MOD>;
int main() {
    Combination<mint> comb(200010);
    int N, M, K; scanf("%d%d%d", &N, &M, &K);
    
    mint ans(0);
    for(int i=0; i<=K; i++) {
        mint a = comb.C(K, i) * mint(M+K-i).pow(N);
        if(i % 2) ans -= a;
        else ans += a;
    }
    ans *= comb.C(M, K);
    printf("%lld\n", ans.v);
    return 0;
}