Pythonでk日後の監獄の独房の状態を求める方法【サイクル検出で高速化】

監獄には8つの独房が一列に並んでおり、それぞれの状態はリスト内の0と1で表されます。1は入居中(占有)、0は空室を意味します。

毎日、次のルールに従って独房の状態が更新されます。

• ある独房の両隣の状態が同じ(両方とも占有、または両方とも空室)であれば、その独房は翌日占有になります。
• それ以外の場合は空室になります。
• 両端の独房には隣接する独房が1つしかないため、常に空室になります。

この記事では、k日後の独房の状態を効率的に求めるPythonプログラムを解説します。

問題の例

たとえば、初期状態が nums = [1, 0, 1, 0, 0, 0, 0, 0] で k = 1 の場合、出力は [0, 1, 1, 0, 1, 1, 1, 0] になります。最初と最後のセルは2つの隣を持てないため、決して占有されない点に注目してください。

アプローチ:サイクル検出による高速化

単純にk日分のシミュレーションを繰り返すことも可能ですが、kが非常に大きい場合(10億日など)は非効率です。

ここで重要なのは、独房の状態は有限個しか存在しないという点です。両端が常に0であるため、中央の6マスの組み合わせは高々2⁶ = 64通りです。したがって、日数を進めていくと必ず同じ状態が再び現れ、サイクル(周期)に入ります。

この周期性を利用すると、以下の手順で効率的に答えを求められます。

アルゴリズムの手順

1. next_day_state()関数を定義する:現在の状態を受け取り、翌日の状態を返す
   • new_cells := cells のコピーを作成
   • new_cells[0] := 0、new_cells[7] := 0(両端は常に空室)
   • j が 1 から 6 の範囲で:
     • cells[j-1] と cells[j+1] が等しければ new_cells[j] := 1
     • そうでなければ new_cells[j] := 0
   • new_cells を返す
2. メイン処理(solveメソッド):
   • seen := 新しい辞書(これまでに見た状態を記録)
   • flag := False、i := 0
   • i < N の間、以下を繰り返す:
     • ns := next_day_state(cells)
     • ns が未登録なら、seen に記録する
     • 既に登録済みなら(= サイクル検出)、flag := True としてループを抜ける
     • cells := ns、i := i + 1
   • flag が True の場合:
     • N := N mod(seen の要素数)—— 周期で割った余りだけ残りのシミュレーションを行う
     • i := 0
     • i < N の間、next_day_state を適用し続ける
   • cells を返す

実装例

import copy
class Solution:
    def next_day_state(self, cells):
        new_cells = copy.copy(cells)
        new_cells[0] = 0
        new_cells[7] = 0
        for j in range(1, 7):
            if cells[j - 1] == cells[j + 1]:
                new_cells[j] = 1
            else:
                new_cells[j] = 0
        return new_cells

    def solve(self, cells, N):
        seen = dict()
        flag, i = False, 0

        while i < N:
            ns = self.next_day_state(cells)
            if tuple(ns) not in seen:
                seen[tuple(ns)] = True
            else:
                flag = True
                break
            cells = ns
            i += 1

        if flag:
            N = N % len(seen)
            i = 0
            while i < N:
                ns = self.next_day_state(cells)
                i += 1
                cells = ns
        return cells

ob = Solution()
nums = [1, 0, 1, 0, 0, 0, 0, 0]
k = 1
print(ob.solve(nums, k))

入力

[1, 0, 1, 0, 0, 0, 0, 0], 1

出力

[0, 1, 1, 0, 1, 1, 1, 0]

計算量について

サイクル検出により、シミュレーションは最長でも周期の長さ(最大64ステップ)+ 余りの日数で完了します。kがどれほど巨大でも、計算時間はほぼ一定に抑えられるため、LeetCodeの「Prison Cells After N Days」のような制約(N ≤ 10⁹)でも十分に対応できます。