Pythonで階段を最後まで登る最小コストを求めるプログラムの実装方法

数値のリスト stairs と整数 k が与えられたとします。現在、私たちは0番目の階段におり、stairs の最後のインデックスにある階段まで登ることが目標です。stairs[i] の値はそのインデックスに到達する際にかかるコストを表し、各ステップで1段から k 段まで任意の段数を一度にジャンプできます。このとき、最後の階段まで登るために必要な最小コストを求めるのがこの問題です。

問題例

例えば、入力が stairs = [4, 11, 11, 3, 2]、k = 3 の場合、出力は 9 になります。これは、コスト [4, 3, 2] の階段を選んで登ることで合計コストを最小化できるためです。

解法のアプローチ

この問題は、動的計画法(DP)とスライディングウィンドウ最小値のテクニック(単調デック)を組み合わせることで効率的に解けます。手順は以下の通りです。

• 両端キュー(deque)q を用意し、ペア (stairs[0], 0)(累積コスト, インデックス)を挿入します。

• i を 1 から stairs のサイズまで繰り返します。

  • i - q[0][1] > k の間、つまり先頭の要素がジャンプ可能な範囲外になったら、q の左端から要素を削除します。

  • curcost := q[0][0] + stairs[i] として、現時点での最小累積コストに現在の階段のコストを加算します。

  • q が空でなく、かつ curcost が q の末尾要素の累積コスト以下である間、末尾の要素を削除します。

  • ペア (curcost, i) を q の末尾に挿入します。

• 最後に、q の末尾要素の累積コスト(最初の値)を返します。

Pythonでの実装例

以下の実装を見ると、理解がより深まるでしょう。

from collections import deque

class Solution:
   def solve(self, stairs, k):
       q = deque([(stairs[0], 0)])
       for i in range(1, len(stairs)):
           while i - q[0][1] > k:
               q.popleft()
           curcost = q[0][0] + stairs[i]
           while q and curcost <= q[-1][0]:
               q.pop()
           q.append((curcost, i))
       return q[-1][0]

ob = Solution()
stairs = [4, 11, 11, 3, 2]
k = 3
print(ob.solve(stairs, k))

入力

[4, 11, 11, 3, 2], 3

出力

9

アルゴリズムのポイント

このアルゴリズムでは、デック内の累積コストが常に単調増加となるように管理しています。これにより、「直近 k 段以内の最小累積コスト」を先頭要素から O(1) で参照でき、各要素の追加・削除はそれぞれ高々1回しか発生しないため、全体の計算量は O(n) に抑えられます。素朴なDPでは O(n × k) かかるところを、この単調デックのテクニックによって大幅に高速化できるのが大きな特徴です。