概要

高橋君が石 \(1\) から石 \(N\) まで最大 \(K\) 個先の石へジャンプして進むとき、踏んだ石のコストの合計の最小値を求める問題です。

考察

まずは動的計画法（DP）を用いて考えます。実装に合わせて石の番号を \(0\)-indexed（\(0\) から \(N-1\)）として扱います。 dp[i] を「石 \(i\) に到達するための最小コスト」と定義すると、石 \(i\) には石 \(i-K, i-K+1, \ldots, i-1\) のいずれかからジャンプしてくることができます。したがって、遷移式は以下のようになります。

\[ dp[i] = \min_{1 \leq j \leq \min(i, K)} (dp[i-j]) + A_i \]

しかし、この式を素朴にループで計算すると、各 \(i\) について最大 \(K\) 回の比較を行うため、全体の時間計算量は \(O(NK)\) となります。制約として \(N\) と \(K\) は最大 \(10^6\) になるため、最悪ケースでは計算回数が \(10^{12}\) オーダーに達し、TLE（実行時間超過）となってしまいます。

これを解決するためには、「長さ \(K\) のスライドする区間内の最小値」を高速に求める必要があります。

アルゴリズム

区間の最小値を高速に求めるために「スライド最小値（Sliding Window Minimum）」という手法を用います。これには両端キュー（deque）を使用します。

キューには「最小値の候補となる石のインデックス」を格納し、以下の性質が常に保たれるように管理します。 - キューの中のインデックスは昇順（古い順） - キューの中のインデックスに対応する dp の値は単調増加（先頭が常に最小値）

石 \(i\) の計算を行う際、具体的に以下の手順でキューを操作します。

1. 古い要素の削除 キューの先頭にあるインデックスが \(i-K\) より小さい（ジャンプできない距離になった）場合、キューの先頭から取り除きます（popleft）。
2. 最小値の取得と更新 この時点で、キューの先頭には必ず区間 \([i-K, i-1]\) における dp の最小値を持つインデックスが入っています。これを使って dp[i] = dp[dq[0]] + A[i] を計算します。
3. 無駄な要素の削除 新しく計算した dp[i] をキューに追加する前に、キューの末尾から「dp[i] 以上の値を持つインデックス」をすべて取り除きます（pop）。なぜなら、これらの要素は dp[i] より値が大きく、しかも dp[i] より古いため、今後区間の最小値になる可能性が絶対にないからです。
4. 新しい要素の追加 キューの末尾に現在のインデックス \(i\) を追加します。

計算量

• 時間計算量: \(O(N)\) 各インデックス \(i\) はキューに対して最大でも1回追加され、1回削除されるだけです。したがって、内側の while ループが実行される合計回数は全体で \(N\) 回以下となり、時間計算量は \(O(NK)\) から \(O(N)\) へと劇的に改善されます。
• 空間計算量: \(O(N)\) 要素数 \(N\) の DP 配列と、最大で要素数 \(N\) となるキューを使用するため、\(O(N)\) となります。

実装のポイント

• 両端キューの利用: Python では collections.deque を使うことで、両端からの要素の追加・削除を \(O(1)\) で高速に行うことができます。リスト（配列）で代用すると先頭の削除に \(O(N)\) かかってしまうため注意が必要です。

• 高速な入出力: 入力される数値が最大 \(10^6\) 個と非常に多いため、標準の input() をループで回すのではなく、sys.stdin.read().split() を用いて入力を一括で読み込むことで、実行時間を大幅に短縮しています。

  ソースコード

import sys
from collections import deque

def main():
    input_data = sys.stdin.read().split()
    if not input_data:
        return
    
    N = int(input_data[0])
    K = int(input_data[1])
    A = [int(x) for x in input_data[2:N+2]]
    
    dp = [0] * N
    dp[0] = A[0]
    
    dq = deque([0])
    
    for i in range(1, N):
        while dq and dq[0] < i - K:
            dq.popleft()
            
        dp[i] = dp[dq[0]] + A[i]
        
        while dq and dp[dq[-1]] >= dp[i]:
            dq.pop()
        dq.append(i)
        
    print(dp[N-1])

if __name__ == '__main__':
    main()

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