AtCoder: 全国統一プログラミング王決定戦本戦 - E問題: Erasure

本戦中にはなんとか解けたけど、通り数計算まわりで無限にハマってつらくなった問題。
こんなハマり方はしないようにしたい...。

問題概要

1から ${N}$ までの番号付きの ${N}$ 個のブロックがある。
${1 \le l \le r \le N, r - l \ge K}$ を満たす ${(l, r)}$ を選び、 ${l}$ 以上 ${r}$ 以下のブロックを消すことを考える。

この時、各 ${(l, r)}$ で消すか消さないかの組み合わせは ${2^{(N-K) (N-K+1)/2}}$ 通り考えられる。
この組み合わせの中で ${N}$ 個のブロック全て消せる組み合わせは何通り存在するか？ ${10 ^ 9 + 7}$ で割った余りで答えよ。

制約

${1 \le N \le 5000}$
${0 \le K \le N-1}$

解法

dpで解く。

以下のような、想定解法よりシンプルなdpを考えてた。

${dp[i] = }$ 番号iのブロックまで全てのブロックを消した時の通り数

このdpでは ${1 \le i \le N}$ について、 ${dp[i]}$ を ${dp[j-1] (1 \le j \le i)}$ から求める。

計算の考え方として、番号jのブロックが残っている時に、ブロックjからブロックiまでを一気に消す場合(つまり、 ${(l, r) = (k, i) (k \le j)}$ で消される場合) の通り数を考えて値を伝搬させる。
${i}$ と ${j}$ の関係に応じて以下の2つのケースに分け、その和を計算することで ${dp[i]}$ を求める。

そして最後まで計算した後、 ${dp[N]}$ を解として出力すればよい。

1. ${i-K < j}$ の場合

f:id:smijake3:20190219231041p:plain — i-K < j の時の通り数の計算

番号jのブロックが残っている時に、ブロックjからブロックiを一気に消すためには、少なくとも
${(1, i), (2, i), ..., (i-K, i)}$
のいずれかで消す必要があり、消せないケースはこれらのいずれも行わなかった場合のみである。
そのため、 ${2^{(i-K)}-1}$ 通りの消し方が存在する。

よって、 ${dp[j-1] \times (2^{(i-K)}-1)}$ を ${dp[i]}$ に加えればよい。

2. ${j \le i-K}$ の場合

f:id:smijake3:20190220010425p:plain — j ≦ i-K の時の通り数の計算

まず、ブロックjからブロックiまでを一気に消すためには、少なくとも
${(1, i), (2, i), ..., (j, i)}$
のいずれかで消す必要があり、これで消せる通り数は ${2^j - 1}$ となる。

また ${j \le i-K}$ の場合では、番号jのブロックを消さずに、番号j+1から番号iまでのブロックを消す通り数も考慮する必要がある。
この通り数は ${j < l, r \le i}$ となる ${(l, r)}$ について考慮した $2^{(i-K-j+1)(i-K-j)/2}$ 通りとなる。

よって、 $dp[j-1] \times (2^j - 1) \times 2^{(i-K-j+1)(i-K-j)/2}$ を ${dp[i]}$ に加えればよい。

実装

計算量は ${O(N^2)}$

提出コード(PyPy3): Submission #4300115 - 全国統一プログラミング王決定戦本戦

N, K = map(int, input().split())
MOD = 10**9 + 7
 
dp = [0]*(N+1)
dp[0] = 1
 
# (2**k % MOD) の O(N^2) の前計算
pw2 = [1]*(N**2+1)
for i in range(N**2):
    pw2[i+1] = pw2[i] * 2 % MOD
 
for i in range(K+1, N+1):
    r = 0
    for j in range(i):
        b = 1
        if j < i-K:
            b = pw2[(i-j-K-1)*(i-j-K)//2] * (pw2[j+1]-1) % MOD
        else:
            b = pw2[i-K] - 1
        r += dp[j] * b % MOD
    dp[i] = r % MOD
 
print(dp[N])