累積和を使う動的計画法

この記事は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は絶対に停まるのには注意してください。
f:id:shindannin:20151224005704p:plain:w700

  • 状態は、
    • 何駅目まで来たか
    • 現在何駅連続で停まっているか
  • 遷移は
    • 次の駅に停まる→連続で停まった回数が+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
f:id:shindannin:20151224005712p:plain:w700
つまり、計算量減らして高速化するには、

  1. 状態数を減らす
  2. 遷移数を減らす

のが基本です。今回は、状態数が多いので、状態数を減らしましょう。1連続停車~K-1連続停車までを1つの状態にします*2
f:id:shindannin:20151224005718p:plain:w700
駅に停まらないケース(青線)は、矢印元と矢印先が同じ状態なので、そのまままとめられます。簡単。
f:id:shindannin:20151224005722p:plain:w700
ただし、駅に停まるケース(赤線)はどうでしょう?1番上のK-1回連続(先の例だと4連続)のケースは、駅にこれ以上停まれないので失われています。この部分は差し引かないとダメですが、状態をまとめてしまってはK-1連続停車の回数がわかりません。この部分をどうやって求めるかが、累積和を使った動的計画法のポイントです。
f:id:shindannin:20151224005727p:plain:w700
最初の図をよく見ると、回数が斜め方向に同じになることがわかります。改めて考えてみると、n駅でK-1連続停車した回数は、n-(K-1)駅で0連続停車した回数といっしょです。
f:id:shindannin:20151224005733p:plain:w700
0連続停車した回数は分かるので、これで計算できます。例えば駅5の1連続~K-1回連続は、4+4-1=7と求められます。
f:id:shindannin:20151224005738p:plain:w700
これで計算量は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;
}

まとめ

動的計画法の高速化には、状態数を減らすのが方法の1つです。累積和を使って減らすことができます。累積和の部分と累積和ではない部分の関係をどう立式するかで、アイディアが必要になってきます。
ただ、今回の方針は一例にすぎず、同じ問題でも多数のやり方があります。ぜひ「Typical DP 準急」で検索して、他の方針との違いも見てみてください。

応用問題として、ちゃっかりこの問題をあげておきます。累積和だけではありませんが、チャレンジしてみてください。
Contest Page | CodeChef


あす24日クリスマスイブは、競技プログラミングアドベントカレンダー創始者のtanzakuさんと、実際に役立った例を紹介してくださるuvarimeronさんです。楽しみですね。それでは、楽しいクリスマスを!

*1:本当は1遷移あたりの計算量も掛けないといけませんが、O(1)の問題のほうが多いでざっくり

*2:着想は、自分の場合は、1連続停車~K-1連続停車まで、だいたい同じような遷移をしているところから思いつきました。