N - Slimes の解説 - めもめも

何の話かと言うと

「区間DP」でうまく解ける問題の例として、こちらを紹介します。

区間DPとは？

「一列にならんだもの（A[1],..., A[N]）を操作する系」の問題で、特定範囲 A[i] 〜 A[j] に制限した場合の答えを dp[i][j] とした上で、幅 d = j - i + 1 を d = 1, 2,... と広げながら dp[ ][ ] を計算していくという手法です。

問題の構成から、自然に区間DPの発想に気がつく例としては、下記の問題があります。

enakai00.hatenablog.com

一方、表題の問題（Slimes）は、事前に「区間DP」の考え方を知らないと、なかなかすぐには解法が思いつかないかもしれません。

愚直な解法

スライムのどこをくっ付けるのかを選択することで、得点と共に状態 s が変化していくので、状態価値関数 dp[s] を用いた解法が考えられます。

s は切れている場所を示した変数だとして、1つの場所 a を繋いだ結果、得点 x が入って、新しい状態 s_n に変化した場合、x + dp[s_n] が「これ以降に得られる合計得点」となります。今回は、合計得点を最小化したいので、すべての a の中で、x + dp[s_n] が最小になるものが選ばれて、これが dp[s] になります。疑似コードだとこんな感じ。

min_cost = infinity
for a in (状態 s において切れている場所):
  s_n = (状態 s から a を繋いだ状態)
  x = (a を繋ぐコスト)
  min_cost = min(min_cost, x + dp[s_n]) 

dp[s] = min_x + dp[s_n]

この遷移により、切れてる場所が少ない状態から、多い状態に向かって計算することができます。もしくは、こちらの最後に触れた、バイナリー値のオーダーで処理するというテクニックも使えます。s を $0\sim 2^{N-1}-1$ の整数として、n ビット目の 0 / 1 が A[n] のすぐ右が「繋がっている／切れている」に対応しているものとします。この場合、ある状態 s において、特定の場所を繋ぐと（特定のビットを 1 から 0 にすると）必ず s の値は減少します。したがって、s を 0 から $2^{N-1}-1$ まで増やしながら遷移計算を進めることができます。

で、これをがんばって実装した例もあるのですが、残念ながら TLE です。

atcoder.jp

この問題の場合、 $1 \le N \le 400$ という前提なので、最悪の場合、 $2^{309}$ という膨大な状態数についてループすることになり、この方法ではとても間に合わないのです。

区間DPによる解法

区間DPは、だいたいの場合は、「後ろを振り返って考える（？）系」になります。今、N 個のスライムのどこか中間部分を切り出した、「d 個のスライムが並んだ状態」があるとします。仮に、A[i] 〜 A[j] を取り出したとしましょう。この状態について問題を解くとして、「（すでに解けているはずの）d-1 個のスライムが並んだ部分列に対する答え」をうまく再利用できないでしょうか？

部分列を使うのですから、A[i] 〜 A[j] のどこかに切れ目を入れて、A[i]〜A[k]、および、A[k+1]A[j] の2種類の問題との関連を考えるんですよね。。。。

あ。

わかった。

一番最後に k の場所をつなぐとすれば、左側を繋ぐコスト dp[i][k] と右側を繋ぐコスト dp[k+1][j]、そして、最後に k の場所で全体を繋ぐコスト sum(A[i:j+1]) がトータルのコストになります。つまり、

dp[i][j] = dp[i][k] + dp[k+1][j] + sum(A[i:j+1])

となるはずです。実際には、つなぐ場所 k は、i 〜 j-1 の範囲で考えられるので、それぞれについて上記のトータルコストを計算して、これらの最小値が dp[i][j] になるはずです。

というわけで、あっさり実装できてしまいました。

import sys

def main(f):
  N = int(f.readline())
  A = [None] + list(map(int, f.readline().split()))

  dp = [[None] * (N+1) for _ in range(N+1)] # dp[i][j] : a_i 〜 a_j を使った場合の答え
  # d = 1
  for i in range(1, N+1):
    dp[i][i] = 0

  for d in range(2, N+1): # 幅dの部分列について計算
    for i in range(1, N+2-d):
      j = i + d - 1
      min_cost = -1
      for k in range(i, j): # 最後に結合する場所で場合分け
        cost = dp[i][k] + dp[k+1][j] + sum(A[i:j+1])
        if min_cost < 0 or cost < min_cost:
          min_cost = cost
      dp[i][j] = min_cost
  print(dp[1][N])

main(sys.stdin)

部分和計算の最適化

ただしこのままでは、まだ TLE します。「部分和計算の最適化」という基本を忘れているためです。（こちらの冒頭で説明しているやつです。）

今の場合は、

        cost = dp[i][k] + dp[k+1][j] + sum(A[i:j+1])

ここに含まれる sum(A[i:j+1]) を S[j] - S[i-1] に置き換えます。（S[i] は事前に計算した A[1] 〜 A[i] の和）

これで無事にパスします。

atcoder.jp