AtCoder DP入門:Frog 1で状態・初期値・遷移を決める

読了 約11分 たびすけ
AtCoder DP入門としてFrog 1の状態・初期値・遷移を示す図

次に読む記事

関連するテーマの記事を、先に確認できます。

Educational DP Contest A – Frog 1は、足場ごとの「そこへ着くまでの最小コスト」を保存する一次元DPで解けます。説明では、公式の足場番号に合わせて f(i) を「足場iへ到達する最小コスト」と定義し、Pythonのコードでは dp[i] を「足場i+1へ到達する最小コスト」として扱います。最初の二つの状態を決めたら、足場の番号が小さい順に、直前一つと二つから来る二つの候補だけを比較します。

この対応を公式サンプルで追い、入力の1始まりをPythonの0始まりへ写す方法、帰納法による正しさ、O(N)時間・O(N)追加メモリ、N=2の境界までを一続きで確認します。最後に、状態の意味を変えずに移動候補だけを K 個へ広げる Frog 2へ進みます。

AtCoder公式のFrog 1問題文公式解説に対応した説明です。

同じ足場へ着く経路を一つの状態にまとめる

カエルは足場1から始まり、現在の足場から一つ先または二つ先へ移動します。移動元と移動先の高さが h なら、支払うコストは高さの差の絶対値です。たとえば足場 i から i+1 へ進むコストは |h_i - h_{i+1}| です。

同じ足場へ複数の経路で到着できるとき、次の移動で使える選択肢は到着した経路によって変わりません。したがって、その足場までのコストが高い経路を残しても、以降の移動で低い経路を逆転することはありません。各足場について最小コストだけを保存すれば、全経路を一つずつ列挙せずに済みます。これが、この問題でDPを使う判断です。

表し方 意味 初期値・更新
f(i) 公式と同じ1始まりで、足場 i へ到達する最小コスト f(1)=0
dp[i] Pythonの0始まりで、足場 i+1 へ到達する最小コスト dp[0]=0

足場 ii > 2)へ最後に移動する直前の足場は、i-1 または i-2 です。そのため、1始まりの漸化式は次のようになります。

f(i) = min(f(i-1) + |h_{i-1} - h_i|, f(i-2) + |h_{i-2} - h_i|)

各候補は「その移動元までの最小コスト」と「最後の一歩のコスト」の和です。その二つのうち小さい方を選ぶので、足場の番号が小さい順に計算すれば、必要な状態が計算済みの状態としてそろいます。

公式サンプル1を f(1)=0 から追跡する

公式サンプル1は、次の高さ列です。

4
10 30 40 20

足場1に最初からいるため、f(1)=0 です。以降は、同じ式を一行ずつ適用します。

状態 その状態へ来る候補 最小コスト
f(1) 開始地点なのでコストは 0 0
f(2) 0 + |10 - 30| = 20 20
f(3) min(20 + |30 - 40|, 0 + |10 - 40|) = min(30, 30) 30
f(4) min(30 + |40 - 20|, 20 + |30 - 20|) = min(50, 30) 30

最後の f(4)=30 が足場4へ着く最小コストで、公式サンプル1の出力 30 と一致します。足場4へは足場3から進む経路だけでなく、足場2から二つ飛ぶ経路も候補に残るため、直前二つの状態を比較する必要があります。

Frog 1の入力、出力、制約を確認する

公式問題では、N 個の足場に高さ h_1, h_2, ..., h_N が与えられます。カエルは足場1から出発し、足場 i から足場 i+1 または i+2 へ移動できます。足場Nへ着くまでに支払うコストの合計を最小化します。

項目 公式問題の内容
入力1行目 足場の個数 N
入力2行目 N 個の高さ h_1, h_2, ..., h_N
制約 2 ≤ N ≤ 1051 ≤ h_i ≤ 104。入力はすべて整数
出力 足場Nへ着くまでの最小コストを一つ出力

問題文の h_1 はPythonでは H[0]h_NH[N-1] です。この添字の対応を固定すると、説明の f(i) とコードの dp[i] を取り違えずに済みます。

0始まりの dp 配列で完全なPythonコードを書く

コードでは、dp[i] を足場 i+1 の状態として保存します。N ≥ 2 が公式制約で保証されるため、足場2に対応する dp[1] を先に初期化できます。足場3以降に対応する dp[i] は、コードの添字で一つ前の dp[i-1] と二つ前の dp[i-2] から更新します。

N = int(input())
H = [int(x) for x in input().split()]

dp = [0] * N
dp[1] = abs(H[0] - H[1])
for i in range(2, N):
    dp[i] = min(
        dp[i - 1] + abs(H[i - 1] - H[i]),
        dp[i - 2] + abs(H[i - 2] - H[i])
    )

print(dp[-1])

range(2, N)i=2 から i=N-1 までを一度ずつ処理します。これは公式の足場3から足場Nに対応します。各回で参照する H[i-1]H[i-2]dp[i-1]dp[i-2] は必ず存在し、最後の dp[N-1] が足場Nの答えです。

掲載コードを公式サンプルで確かめる

掲載コードへ公式サンプルを入力形式のまま与えると、次の出力になります。サンプル1の中間値は公式の入力と漸化式から計算した追跡です。

サンプル 入力 出力
1
4
10 30 40 20
30
2
2
10 10
0
3
6
30 10 60 10 60 50
40

なぜこの更新で最小コストになるのか

コードの dp[i] が、足場 i+1 へ到達する最小コストを表すことを、足場の順番に沿った帰納法で確認します。

まず i=0 では、カエルは足場1に最初からいるので、到達コストは dp[0]=0 です。i=1 では、足場2へは足場1から直接移動するしかありません。そのコストは abs(H[0] - H[1]) なので、初期化した dp[1] も正しい値です。

次に、i ≥ 2 について、dp[0] から dp[i-1] までの状態が正しいと仮定します。足場 i+1 へ入る最後の移動は、足場 i からの一歩か、足場 i-1 からの二歩です。仮定により、足場 i までの最小コスト dp[i-1] と、足場 i-1 までの最小コスト dp[i-2] は正しい値です。したがって、最後の一歩を足した二つの候補 dp[i-1] + abs(H[i-1] - H[i])dp[i-2] + abs(H[i-2] - H[i]) はそれぞれの移動元からの最小コストであり、その最小値を保存する dp[i] も足場 i+1 へ着く最小コストになります。

この議論を足場Nまで繰り返せるため、最後に出力する dp[-1] は求める最小コストです。計算順を左から右にしているのは、各更新の時点で dp[i-1]dp[i-2] をすでに確定させるためです。

計算量と境界を確認する

i=2 から N-1 までの各状態を一度だけ計算し、各回で二つの候補を比較します。そのため時間計算量は O(N) です。dp 配列に N 個の値を保存するので、追加メモリは O(N) です。

  • N=2dp[1] を初期化した後、range(2, N) は空になります。ループを回さず、足場1から足場2へのコストをそのまま出力します。
  • 同じ高さ:高さの差が0なので、その移動の候補コストも0です。特別扱いせず、abs の計算にそのまま含めます。
  • 添字の端:ループの最初は i=2 なので i-2=0 です。最後は i=N-1 なので、参照先は配列の最後から一つ前と二つ前に収まり、H[i]H[N-1] になります。

N=1、空入力、障害物、円環状の移動は公式の制約・問題文の対象外です。このコードの境界として、それらを追加の仕様とは扱いません。

同じ状態のまま Frog 2で候補範囲を広げる

次の公式問題 Educational DP Contest B – Frog 2では、足場 i から i+1 だけでなく、i+K まで移動できます。状態の意味はFrog 1と同じく、f(i) を「足場 i へ到達する最小コスト」とします。変えるのは、最後の移動元として調べる範囲だけです。

1始まりで足場 i の候補を表すと、max(i-K, 1) ≤ j ≤ i-1 にある足場 j からのコストを調べ、次の最小値を取ります。

f(i) = min(f(j) + |h_j - h_i|)max(i-K, 1) ≤ j ≤ i-1

Frog 1では候補が直前二つに固定されていたため各状態を定数時間で更新できました。Frog 2では最大 K 個の候補を確認するので、公式解説の計算量は O(NK) です。Pythonの0始まりでは候補範囲を range(max(i-K, 0), i) と書きます。

練習では、まずFrog 1と同じ状態の意味を一行で書き、ある i に対して候補となる j の範囲を列挙します。その後、Frog 1の二候補がFrog 2で直前 K 個以内へ広がったこと、入力に K が加わったこと、計算量が O(N) から O(NK) へ変わることを公式サンプルで確かめると、状態定義と遷移範囲の役割を分けて考えられます。

Frog 2の公式問題公式解説を使い、まず候補範囲を書き出してからコードに進んでください。Frog 1で確認したのは、一つの状態が何を表すかを固定し、計算済みの前の状態から次を作る流れです。Frog 2では、その流れを保ったまま最後の一歩の候補だけを増やします。