Pythonで石を渡って川を越えられるか判定するプログラム

ソートされた整数リスト stones が与えられ、これは渡ろうとしている川に置かれた石の位置を表しているとします。川を渡り切るためには、必ず最後の石までたどり着かなければなりません。各ステップでは、直前のジャンプ距離を k としたとき、(k − 1, k, k + 1) のいずれかの距離だけ前へジャンプすることができます。この条件のもとで、川を渡ることが可能かどうかを判定するのが本問題です。

問題の例

たとえば入力が stones = [0, 1, 3, 4, 5, 6, 8, 9, 13] の場合、答えは True になります。位置 0 からスタートし、まず 1 単位ジャンプして石 1 へ、次に 2 単位ジャンプして石 3 へ、さらに 2 単位で石 5 へ、続いて 3 単位で石 8 へ、最後に 5 単位ジャンプして最後の石である 13 に到達できるからです。

解法のアプローチ

この問題は、深さ優先探索（DFS）によるバックトラッキングで効率よく解けます。手順は以下の通りです。

• start := A[0]、end := A の最後の要素とする
• A を一意な要素のみを持つ集合（set）に変換し、位置の存在確認を高速に行えるようにする
• 関数 check() を定義する。初期値は pos := start、prev := 0
• pos が end と等しければ True を返す
• [prev − 1, prev, prev + 1] の各 jump に対して以下を試す
  • jump >= 1 であれば、next_pos := jump + pos を計算する
  • next_pos が A に存在し、かつ check(next_pos, jump) が真であれば True を返す
• どのジャンプでもゴールに届かなければ False を返す
• メイン処理から check() を呼び出し、その結果を返す

重要なポイントは、ジャンプできる距離が直前のジャンプ距離に依存するため、現在位置だけでなく直前のジャンプ距離も再帰の状態として引き継ぐ必要がある点です。また、ジャンプ距離は最低 1 以上でなければならないため、jump >= 1 のチェックを忘れないようにしましょう。

Python 実装例

それでは、実際のコードを見て理解を深めましょう。

class Solution:
    def solve(self, A):
        start, end = A[0], A[-1]
        A = set(A)
        def check(pos=start, prev=0):
            if pos == end:
                return True
            for jump in [prev - 1, prev, prev + 1]:
                if jump >= 1:
                    next_pos = jump + pos
                    if next_pos in A and check(next_pos, jump):
                        return True
            return False
        return check()

ob = Solution()
stones = [0, 1, 3, 4, 5, 6, 8, 9, 13]
print(ob.solve(stones))

入力

[0, 1, 3, 4, 5, 6, 8, 9, 13]

出力

True

計算量と改善のヒント

この実装では、「現在位置」と「直前のジャンプ距離」の組み合わせを状態として再帰的に探索します。石の数を n とすると、最悪の場合の時間計算量は O(n²) 程度になります。より大規模な入力に対応するには、一度訪れた (位置, ジャンプ距離) の組み合わせを辞書などでキャッシュするメモ化を追加すると、無駄な再探索を排除でき、パフォーマンスが大きく向上します。コーディング面接や競技プログラミングでは、この最適化まで言及できると好印象です。