動的計画法を用いた組み合わせ計算

競プロ初心者ですが、アルゴリズムの勉強も兼ねて競プロ向けC++ライブラリを整備していくことにしました。 C++ライブラリと言っていますが、Pythonも追加するかもしれません。 github.com 今回はタイトルの通り、DP(Dynamic Programming, 動的計画法)を用いた組み合わせ計算を実装しました。

アルゴリズム

漸化式

n個からk個を選ぶ組み合わせの数をcomb(n, k)とすると、下記の漸化式が成り立ちます。

comb(n, k) = comb(n - 1, k - 1) + comb(n - 1, k)

上記漸化式についての説明は、数学ガール/乱択アルゴリズムの説明が大変分かりやすいです。持っていない人は買いましょう。

数学ガール/乱択アルゴリズム (数学ガールシリーズ 4)

制約条件

実装のことも考えて、下記を制約条件とします。

• n, kは整数
• 0 ≦ k ≦ n
• comb(n, 0) = 1
• comb(0, k) = 0

DP

DPしていきます。変数i, j (0 ≦ i ≦ n, 0 ≦ j ≦ k)について二重ループを回して下記を計算していきます。計算量はO(n*k)になります。

dp[i + 1][j + 1] = (dp[i][j] + dp[i][j + 1]) % MOD;

ここでは競プロらしく1e9+7で剰余を求め、int型に収まるよう計算しています。

また、初期化に関しても同一ループ中で実行しています。

if (j == 0) dp[i][j] = 1;
else if (i == 0) dp[i][j] = 0;

ソースコード

https://github.com/aoooooooooi/competitive_programming/blob/master/combinatorics/combination.cpp

#include <iostream>

using namespace std;

const int MAX = 100;
const int MOD = 1e9 + 7;
int dp[MAX][MAX];

int main() {
    int n, k;
    cin >> n >> k;

    for (int i = 0; i + 1 <= n; ++i) {
        for (int j = 0; j + 1 <= k; ++j) {
            if (j == 0) dp[i][j] = 1;
            else if (i == 0) dp[i][j] = 0;
            dp[i + 1][j + 1] = (dp[i][j] + dp[i][j + 1]) % MOD;
        }
    }
    cout << dp[n][k] << endl;
}

おわり

以上です。DPはまだ慣れないですが、漸化式を立てるのはもちろん、初期値の設定が重要だと感じました。