hogecoder

tsutaj 競技プログラミングの記録

動的計画法における空間計算量削減テク

競プロその他

この記事は 「競プロ!!」競技プログラミング Advent Calendar 2017 の 9 日目の記事として書かれました。

競技プログラミングをそれなりにやってきた方となれば、動的計画法に触れたことも多いでしょう。今回は、その動的計画法における空間計算量削減テクを色々紹介したいと思います。

普通のナップザック問題

題材として、ごく普通のナップザック問題を考えてみましょう。

容量 W のナップザックに、品物をいくつか入れたいと考えています。品物は全部で N 個あり、 i 番目の品物は価値 v_i 、容量 w_i です。品物の在庫はそれぞれ 1 個ずつなので、同じ商品を複数個入れることはできません。ナップザックの容量を超えることなく商品を入れる時、入れた商品の価値の総和の最大値を求めてください。

制約は

  • 1 \leq N \leq 10^{4}
  • 1 \leq W \leq 10^{4}
  • 1 \leq v_i \leq 10^{4}
  • 1 \leq w_i \leq 10^{4}
  • 時間制限 3 sec
  • メモリ制限 64 MB

とします。せっかく空間計算量削減の話をするんですから、これくらい厳しくしないとですよね！

まずは普通に実装してみます。結論から言うと、このコードは空間計算量 O(NW) で配列とりすぎなので Memory Limit Exceeded になるはずです。

#include <iostream>
#include <algorithm>
using namespace std;
 
int v[10000], w[10000];
int dp[10001][10001];
int main(){
    int N, W, ans = 0; cin >> N >> W;
    for(int i=0; i<N; i++) cin >> v[i] >> w[i];
    for(int i=0; i<N; i++) {
        for(int j=0; j<=W; j++) {
            dp[i+1][j] = max(dp[i+1][j], dp[i][j]);
            if(j-w[i] >= 0) dp[i+1][j] = max(dp[i+1][j], dp[i][j-w[i]] + v[i]);
            ans = max(ans, dp[i+1][j]);
        }
    }
    cout << ans << endl;
    return 0;
}

このメモリ制限に対応するためには、何か工夫が必要です。どうすれば良いでしょうか？

インデックス \bmod 2 DP

ここで、 DP の遷移に着目します。(該当箇所再掲)

dp[i+1][j] = max(dp[i+1][j], dp[i][j]);
if(j-w[i] >= 0) dp[i+1][j] = max(dp[i+1][j], dp[i][j-w[i]] + v[i]);

i+1 番目の状態を得るために参照している状態は、 i 番目の状態のみです。もっと言うと、 i-1 番目以前はもう使いません。これは無駄そうです。更新に必要な部分だけを保存していくように変えたいですね。

次の状態を得るには今の状態だけ分かっていればいいのですから、 i のインデックスは 2 つあれば十分です。 i の偶奇を気にすると、次のようにすれば良さそうです。

このように、 i の偶奇に応じて参照先を変え、最新 2 つの状態を持っておくことで空間計算量を大幅に削減できます。更新時にどのインデックスを見るかに着目して、必要なものだけ持つ、ということが本質です。この工夫をすることによって、空間計算量は O(W) になりました。

改善後のコードは以下の通りです。

#include <iostream>
#include <algorithm>
using namespace std;
  
int v[10000], w[10000];
int dp[2][10001];
int main(){
    int N, W, ans = 0; cin >> N >> W;
    for(int i=0; i<N; i++) cin >> v[i] >> w[i];
    for(int i=0; i<N; i++) {
        int mo = i % 2;
        for(int j=0; j<=W; j++) {
            dp[mo^1][j] = max(dp[mo^1][j], dp[mo][j]);
            if(j-w[i] >= 0) dp[mo^1][j] = max(dp[mo^1][j], dp[mo][j-w[i]] + v[i]);
            ans = max(ans, dp[mo^1][j]);
        }
        for(int j=0; j<=W; j++) dp[mo][j] = 0;
    }
    cout << ans << endl;
    return 0;
}

dp 配列の 1 次元目が大幅に減りました。これで、メモリ制限に対応できますね。

この実装で注意すべきは、使い終わったインデックスは初期化する ということです。上の実装でいう for(int j=0; j<=W; j++) dp[mo][j] = 0; の部分にあたります。これをやらないと次の更新をするときに本来いらないはずの値が既に入ってしまっており、不具合が生じるからです。

更新の順番工夫 DP

実はもっと空間計算量を削減できます。今まで 2 次元の DP で解いてきましたが、今から紹介する工夫をすることで 1 次元の DP になります。とはいえ、今の方針を無理やり 1 次元にしてしまうと、既に更新した状態を次の更新時に使ってしまい、二重更新が発生します。どうやったらうまくできるでしょうか？

DP 配列の更新を (更新に使われる) → (更新に使う) の矢印で考えると、次のように 左から右に 矢印が引かれると思います。これは、値の更新時に今よりも容量が小さい状態を見るからです。

この性質は、容量の大きい状態から順に確定していくことを意味します。これを利用して、容量の大きい順に DP 配列を更新していくと、品物のインデックスが不要になり、 1 次元の DP で解くことができます。

#include <iostream>
#include <algorithm>
using namespace std;
 
int v[10000], w[10000];
int dp[10001];
int main(){
    int N, W, ans = 0; cin >> N >> W;
    for(int i=0; i<N; i++) cin >> v[i] >> w[i];
    for(int i=0; i<N; i++) {
        for(int j=W; j>=0; j--) {
            if(j-w[i] >= 0) dp[j] = max(dp[j], dp[j-w[i]] + v[i]);
            ans = max(ans, dp[j]);
        }
    }
    cout << ans << endl;
    return 0;
}

逆順ループによる空間計算量削減テクは DP の更新を矢印でイメージすると、適用できるか否かがわかりやすいです。たとえば、XOR を取る DP などはこのように 1 次元 DP に落としこむことはできません。なぜなら、 XOR を取ることで今の数字から増える場合と減る場合の両方があり、更新の矢印の方向に規則性がなく、一度更新したところを再度見てしまう可能性があるからです。

もしこの方法が苦手であっても、大体の DP は先ほどのインデックス \bmod 2 DP で対応できますので、そちらをマスターするのでも十分だと思います。

演習

最後に、今回紹介した空間計算量削減テクが (おそらく) 必須な問題を挙げました。どれも少し難しめですが、ぜひやってみてください。慣れないうちは、ナップザックなどの簡単な DP で感覚を身につけるのもオススメです。

最近はメモリ制限 256 MB の問題が多いため、空間計算量を意識する機会はそんなに多くない印象です。ですが、昔の JOI などはメモリ制限 64 MB の問題が割とあり、意識しなくてはいけない場合があります。覚えて損はないので、この記事を読んで使えるようになってくれたらうれしいです。

明日の Advent Calendar は tozangezan さんです！