読者です 読者をやめる 読者になる 読者になる

Typical DP Contest F-準急

http://tdpc.contest.atcoder.jp/tasks/tdpc_semiexp
 Atcoder過去問の中でも特に埋める楽しみがあると(自分の中で)評判のコンテスト。
 入力は2以上1000000以下の整数NとK。N個の駅があり、1番目とN番目には必ず止まります。また、連続するK個以上の駅に止まってはいけません。このとき停車駅の組み合わせが何通りあるかmod1000000007で求めろ、という問題。
 まず、次のような動的計画法を考えました。
 dp[i][j] := i番目の駅を過ぎて、(i番目の駅を含めて)連続するj個の駅に停車中のときの場合の数
これを実装するとこんな感じ。

const int MOD = 1000000007;
int dp[100][100];
int main(){
    int n, k;
    scanf("%d%d", &n, &k);
    memset(dp, 0, sizeof(dp));
    dp[1][1] = 1;
    REP(i, 2, n + 1){
        if(i != n) rep(j, k) dp[i][0] += dp[i - 1][j];
        REP(j, 1, k) dp[i][j] += dp[i - 1][j - 1];
    }
    int res = 0;
    rep(i, k) res += dp[n][i];
    printf("%d\n", res);
    return 0;
}

 いちおうこれでサンプルは通りますが、nとkが大きいのでメモリが足りません。そこで、さきほどのdp[i][j]の[i]の部分を使いまわすことにします。

int dp[2][1000010];
int main(){
    int n, k;
    scanf("%d%d", &n, &k);
    memset(dp, 0, sizeof(dp));
    int crr = 0, nxt = 1;
    dp[crr][1] = 1;
    REP(i, 2, n + 1){
        memset(dp[nxt], 0, sizeof(dp[nxt]));
        if(i != n) rep(j, k) (dp[nxt][0] += dp[crr][j]) %= MOD;
        REP(j, 1, k) (dp[nxt][j] = dp[crr][j - 1]) %= MOD;
        swap(crr, nxt);
    }
    int res = 0;
    rep(i, k) (res += dp[crr][i]) %= MOD;
    printf("%d\n", res);
    return 0;
}

 これを提出したらTLEになってしまいました。オーダーはn * kとかなので大丈夫かと思ったのですがもうひと工夫必要みたいです。
 大きいループの中を見てみると、dp[crr][j]の合計を求めてdp[nxt][0]とする操作(…①)、dp[crr]をひとつずらしてdp[nxt]とする操作(…②)をしており、これはループを使わなくても書けそうです。
 現在のdpテーブルは
dp[i][0] = dp[i - 1][0] + ... + dp[i - 1][k - 1]
dp[i][j](1 <= j < k) = dp[i - 1][j - 1]
です。新しく、
 dp2[i] = dp[i][0] + ... + dp[i][k - 1] とします。
 このとき、
dp[i][0] = dp[i - 1][0] + ... + dp[i - 1][k - 1]
= dp2[i - 1]
また、
 dp[i][1] + ... + dp[i][k - 1] = dp[i - 1][0] + dp[i - 1][1] + ... + dp[i - 1][k - 2]
                  = dp2[i - 1] - dp[i - 1][k - 1]
ですから、
dp2[i] = dp[i][0] + dp[i][1] + ... + dp[i][k - 1]
= dp2[i - 1] + dp2[i - 1] - dp[i - 1][k - 1]
さらに
dp[i - 1][k - 1] = dp[i - 2][k - 2]
= ...
= dp[i - k][0]
= dp2[i - k - 1]
となり、ようやく
dp2[i] = dp2[i - 1] * 2 - dp2[i - k - 1]
が得られました。
 ところで、k + 1回前の値を引いているわけなのですが、連続してk回以上止まってはだめというルールなのでk回目でも場合分けが必要なはずです。さらに最後は必ず止まるのでここも場合分けが必要で、このへんが少しやっかいですね。

const int MOD = 1000000007;
int dp[1000010];
int main(){
    int n, k;
    scanf("%d%d", &n, &k);
    memset(dp, 0, sizeof(dp));
    dp[0] = 1;
    REP(i, 1, n){
        dp[i] = dp[i - 1] * 2 % MOD;
        if(i >= k + 1) dp[i] = (dp[i] + MOD - dp[i - k - 1]) % MOD;
        if(i == k - 1) dp[i] = (dp[i] + MOD - 1) % MOD;
    }
    printf("%d\n", (dp[n - 1] + MOD - dp[n - 2]) % MOD);
    return 0;
}

 最終的にはとても簡単なプログラムになりました。一発で立式するのも無理ではないと思うのですが(駅に止まるか止まらないかの二択なので二倍、というのは簡単ですし)、まだまだ修練が足りませんね。

追記:queueを使うと簡単ですね。

const int MOD = 1000000007;
int main(){
    int n, k;
    scanf("%d%d", &n, &k);
    queue<int> que;
    int sum = 1;
    que.push(1); que.push(0);
    REP(i, 2, n){
        que.push(sum);
        sum = (sum * 2) % MOD;
        if(que.size() > k){
            sum = (sum + MOD - que.front()) % MOD;
            que.pop();
        }
    }
    printf("%d\n", (sum + MOD - que.front()) % MOD);
    return 0;
}