D - 花壇の区間分け / Dividing the Flower Bed into Sections 解説 by admin
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)\) となります。
- 長さ \(N+1\) の1次元配列
実装のポイント
逆順ループによる \(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 によって生成されました。