累積和を使う動的計画法
この記事はCompetitve Programming Advent Calendar 2015の23日目の記事です。tanzakuさんに感謝
www.adventar.org
今回は、累積和を使う動的計画法についてです。TopCoderのDiv2上位ぐらいの人向けの難易度です。
問題
AtCoder Typical DP Conestの問題です。
F: 準急 - Typical DP Contest | AtCoder
Problem Statement
ある路線には駅 1 から駅 N までの N 個の駅がある。すぬけ君は、この路線に準急を走らせることにした。
準急は、駅 1 に止まり、{駅 2, ..., 駅 N-1} の部分集合に止まり、駅 N に止まる。
連続する K個以上の駅に止まると、客が飽きてしまうので、そのようなことはしない。
準急の停車駅の組み合わせとして何通り考えられるか、mod 1,000,000,007 で求めよ。
Constraints
2≤K≤N≤1000000
最初の方針（ただし、計算時間が間に合わない）
まず答えがでる実例N=9, K=5で手計算すると、以下のように112通りになります。図は矢印にそって足していくだけです。駅1と駅9は絶対に停まるのには注意してください。
- 状態は、
- 何駅目まで来たか
- 現在何駅連続で停まっているか
- 遷移は
- 次の駅に停まる→連続で停まった回数が+1
- 次の駅に停まらない→連続で停まった回数が0に戻る
となるので、動的計画法で
dp[何駅目まで来たか][現在何駅連続で停まっているか]
とすれば、コードは以下のようになります。
#include <iostream> using namespace std; long long dp[123][123]; int main() { int N,K; cin >> N >> K; long long MOD = 1000000007; dp[1][1]=1; for (int n = 1; n <= N-1; ++n) { for(int k = 0; k <= K-1; ++k) { dp[n+1][k+1] += dp[n][k]; // 駅にとまる dp[n+1][k+1] %= MOD; dp[n+1][0] += dp[n][k]; // 駅にとまらない dp[n+1][0] %= MOD; } } long long ans = 0; for(int k = 1; k <= K-1; ++k ) { ans += dp[N][k]; ans %= MOD; } cout << ans << endl; return 0; }
しかし、この問題ではNもKが1000000まで取りえます。計算時間もメモリも間に合いません。
累積和で状態数を減らす
動的計画法の計算量は、ざっくりいうと、以下のように状態数×（１状態あたりの）遷移数になります*1。
つまり、計算量減らして高速化するには、
- 状態数を減らす
- 遷移数を減らす
のが基本です。今回は、状態数が多いので、状態数を減らしましょう。1連続停車～K-1連続停車までを1つの状態にします*2
駅に停まらないケース（青線）は、矢印元と矢印先が同じ状態なので、そのまままとめられます。簡単。
ただし、駅に停まるケース（赤線）はどうでしょう？１番上のK-1回連続（先の例だと4連続）のケースは、駅にこれ以上停まれないので失われています。この部分は差し引かないとダメですが、状態をまとめてしまってはK-1連続停車の回数がわかりません。この部分をどうやって求めるかが、累積和を使った動的計画法のポイントです。
最初の図をよく見ると、回数が斜め方向に同じになることがわかります。改めて考えてみると、n駅でK-1連続停車した回数は、n-(K-1)駅で0連続停車した回数といっしょです。
0連続停車した回数は分かるので、これで計算できます。例えば駅5の1連続～K-1回連続は、4+4-1=7と求められます。
これで計算量はO(N)まで落とせます。コードは以下のようになります。
#include <iostream> using namespace std; long long dp[1234567]; // dp[n] = n駅まできて、現在0連続停車 long long sum[1234567]; // sum[n] = n駅まできて、現在1～K-1連続停車 int main() { int N,K; cin >> N >> K; long long MOD = 1000000007; dp[0]=1; sum[1]=1; for (int n = 1; n <= N-1; ++n) { dp[n+1]=sum[n]+dp[n]; dp[n+1]%=MOD; sum[n+1]=sum[n]+dp[n]; sum[n+1] %= MOD; if(n+1>=K) { sum[n+1]=sum[n+1]-dp[n+1-K]+MOD; sum[n+1] %= MOD; } } cout << sum[N] << endl; return 0; }
まとめ
動的計画法の高速化には、状態数を減らすのが方法の１つです。累積和を使って減らすことができます。累積和の部分と累積和ではない部分の関係をどう立式するかで、アイディアが必要になってきます。
ただ、今回の方針は一例にすぎず、同じ問題でも多数のやり方があります。ぜひ「Typical DP 準急」で検索して、他の方針との違いも見てみてください。
応用問題として、ちゃっかりこの問題をあげておきます。累積和だけではありませんが、チャレンジしてみてください。
Contest Page | CodeChef
あす24日クリスマスイブは、競技プログラミングアドベントカレンダー創始者のtanzakuさんと、実際に役立った例を紹介してくださるuvarimeronさんです。楽しみですね。それでは、楽しいクリスマスを！
