ABC132 D問題：Blue and Red Balls - 忘れても大丈夫

問題. Blue and Red Balls

$K$ 個の青いボールと $N - K$ 個の赤いボールがある．この $N$ 個のボールを一列に並べたものを $b$ としたとき， $b$ から青いボールが無くなるまで次の操作を行う． $b$ の中にある連続する青いボールの部分列を任意に取り除き，取り除いてできた列を $b$ とする． $b$ に対する任意の操作の中で操作回数の最小値を $b$ の最小操作回数と呼ぶ．任意の $1 \le i \le K$ に対して最小操作回数が $i$ となるボールの並べ方が何通りあるかを求めよ．

制約： $1 \le K \le N \le 2,000$

解法

最小操作回数が $i$ となるボールの並べ方が何通りあるかを考える．これは， $K$ 個のボールを空でない $i$ 個に分割する組合せの数に等しく重複組合せで求めることができる．ここ（重複組合せの公式と例題（玉，整数解の個数） | 高校数学の美しい物語）の正の整数解の個数を求める問題に等しく， $x_1 + x_2 + \cdots + x_i = K$ という方程式の正の整数解の個数で $\binom{K - 1}{i - 1}$ 通りとなる．
次に，青いボールが空でない $i$ 個に分割されたときの，赤いボールの並べ方を考える．赤いボールは青いボールを分割するために，青いボールを分割する $i - 1$ 箇所に少なくとも 1 つの赤いボールがなければならない． $i - 1$ 箇所に 1 個ずつ赤いボールがあるとき，青いボールは $i$ 個に分割されているので，赤いボールの残りの $N - K - i + 1$ 個は $i + 1$ 箇所の中で好きな所に配置することができる．これは， $x_1 + x_2 + \cdots + x_{i + 1} = N - K - i + 1$ という方程式の非負の整数解の個数で $\binom{N - K + 1}{i}$ 通りとなる．
以上から最小回数が $i$ となるボールの並べ方は $\binom{K - 1}{i - 1} \times \binom{N - K + 1}{i}$ 通りとなる．二項係数はパスカルの三角形で求めても間に合う．

計算時間： $O(N^2)$

ソースコードを表示

#include <bits/stdc++.h>

using namespace std;

using ll = long long;

int main() {
    cin.tie(0); ios::sync_with_stdio(false);
    // cout << setprecision(8) << setiosflags(ios::fixed);

    const ll MOD = 1e9 + 7;

    ll n, k;
    cin >> n >> k;

    vector<vector<ll>> comb(n + 1, vector<ll>(n + 1));
    for (int i = 0; i <= n; ++i) comb[i][0] = 1;
    for (int i = 1; i <= n; ++i)
        for (int j = 1; j <= i; ++j)
            (comb[i][j] += comb[i - 1][j - 1] + comb[i - 1][j]) %= MOD;

    for (ll i = 1; i <= k; ++i)
        cout << (comb[k - 1][i - 1] * comb[n - k + 1][i]) % MOD << '\n';

    return 0;
}

方針はすぐ立ったけど実装であたふたした