Official

D - 花壇の区間分け / Dividing the Flower Bed into Sections Editorial by admin Editorial for printing

Gemini 3.1 Pro (Thinking)

概要

\(N\) 個の要素からなる数列を、順番を保ったまま \(K\) 個の連続した区間に分割する際、各区間の「最大値と最小値の差」の合計を最小化する問題です。

考察

連続する区間に分割してコストの最適化を行う問題は、動的計画法(DP)と非常に相性が良いです。

「先頭から \(i\) 番目までの花を \(j\) 個の区間に分けるときの最小コスト」を状態として持てば解けそうだと予想できます。このとき、最後の \(j\) 番目の区間を「\(p+1\) 番目から \(i\) 番目まで」とすると、その区間のコスト(最大値と最小値の差)を足し合わせることで遷移ができます。

素朴なアプローチとして、考えられるすべての区間 \([p+1, i]\) について毎回最大値と最小値を計算し直すと、1回の遷移に \(O(N)\) かかり、全体の計算量が \(O(K N^3)\) となってしまいます。Pythonなどの言語では実行時間制限(TLE)に引っかかる可能性があります。

これを解決するためには、区間の開始位置である \(p\) を \(i-1\) から左へ向かって順番に動かしていくのが有効です。要素が1つずつ追加されていく形になるため、最大値と最小値を \(O(1)\) で更新しながら計算でき、全体の計算量を \(O(K N^2)\) に落とすことができます。

アルゴリズム

動的計画法(DP)を用いて解を進めます。 本来は2次元配列 dp[j][i] を考えますが、実装上はメモリを節約するために1次元配列 dp[i] を使い、\(j\) のループごとに配列を更新します。

dp[i] を「直前のステップ(\(j-1\) 個の区間に分割)において、先頭から \(i\) 本の花を分割したときの総手間の最小値」と定義します。

遷移式は以下のようになります。 $\( \text{new\_dp}[i] = \min_{j-1 \le p < i} \{ \text{dp}[p] + (\text{区間 } [p+1, i] \text{ の最大値} - \text{最小値}) \} \)$

具体的な手順: 1. 初期化: dp[0] = 0 とし、それ以外は十分に大きな値(INF)とします。 2. \(j = 1, 2, \dots, K\) (区間の数)について、以下の更新を行います。 3. 新しい配列 new_dp を INF で初期化します。 4. \(i = j, j+1, \dots, N\) (分割する花の数)について調べます。 5. 最後の区間の境界 \(p\) を \(i-1\) から \(j-1\) まで逆順(右から左)にループさせます。 - このとき、区間 \([p+1, i]\) に含まれる花の高さの最大値 max_val と最小値 min_val を更新します。 - dp[p] + max_val - min_val を計算し、これまでの最小値より小さければ new_dp[i] を更新します。 6. dp を new_dp に置き換え、次の \(j\) へ進みます。 7. 最終的に dp[N] が求める最小の総手間となります。

計算量

  • 時間計算量: \(O(K N^2)\)
    • 状態数が \(O(NK)\) 個あり、各状態の遷移において \(p\) を最大 \(N\) 回ループします。ループ内で最大値・最小値の更新は \(O(1)\) で行えるため、全体の時間計算量は \(O(K N^2)\) です。制約の \(N, K \le 200\) において、最大でも \(200 \times 200 \times 200 \approx 8 \times 10^6\) 回の演算となり、十分に高速です。
  • 空間計算量: \(O(N)\)
    • 長さ \(N+1\) の1次元配列 dp と new_dp のみを保持するため、\(O(N)\) となります。

実装のポイント

  • 逆順ループによる \(O(1)\) 更新: \(p\) を \(i-1\) から小さくする方向へ動かすのが最大のポイントです。これにより、「右端が固定され、左端が1つずつ伸びていく区間」の最大値・最小値を調べることになり、毎回要素を1つ比較するだけで min_val と max_val を更新できます。

  • 配列の使い回し(インプレースなDP): 2次元配列 dp[j][i] を確保しても空間計算量は \(O(NK)\) で問題ありませんが、直前の \(j-1\) の結果しか使わないため、1次元配列を2つ(dp と new_dp)用意して使い回すことで、より省メモリでキャッシュ効率の良い実装になります。

    ソースコード

import sys

def main():
    input = sys.stdin.read
    data = input().split()
    if not data:
        return
    N = int(data[0])
    K = int(data[1])
    A = [int(x) for x in data[2:N+2]]
    
    INF = 10**18
    dp = [INF] * (N + 1)
    dp[0] = 0
    
    for j in range(1, K + 1):
        new_dp = [INF] * (N + 1)
        for i in range(j, N + 1):
            min_val = A[i-1]
            max_val = A[i-1]
            res = INF
            for p in range(i - 1, j - 2, -1):
                val = A[p]
                if val < min_val:
                    min_val = val
                elif val > max_val:
                    max_val = val
                
                cost = dp[p] + max_val - min_val
                if cost < res:
                    res = cost
            new_dp[i] = res
        dp = new_dp
        
    print(dp[N])

if __name__ == '__main__':
    main()

この解説は gemini-3.1-pro-thinking によって生成されました。

posted:
last update: