Pythonで解く「鍵と部屋」問題：BFSで全ての部屋を訪問できるか判定する方法

N個の部屋があり、部屋0からスタートするとします。各部屋には0、1、2、…、N-1の中から一意な番号が割り当てられており、それぞれの部屋には次の部屋へ入るための鍵がいくつか置かれていることがあります。言い換えると、各部屋iにはrooms[i]という鍵のリストが存在し、各鍵rooms[i][j]は[0, 1, …, N-1]の範囲の整数です(ここでNは部屋の総数)。鍵rooms[i][j] = vは、番号vの部屋を開けられることを意味します。

例えば、入力が[[1], [2], [3], []]であれば、出力はtrueになります。

問題を解く前に押さえておきたいポイント

• 最初、すべての部屋は施錠されています(ただし部屋0を除く)。
• 部屋間は自由に行き来できます。
• すべての部屋に入室できる場合にのみ、trueを返す必要があります。

具体的には、部屋0からスタートして鍵1を手に入れ、部屋1へ移動して鍵2を取得します。さらに部屋2から鍵3を取り、部屋3を訪問します。このようにしてすべての部屋を訪問できれば、trueを返します。

アルゴリズムの手順

この問題は、グラフの幅優先探索(BFS)として捉えることで効率的に解けます。手順は以下の通りです。

• 空のキューを作成し、全部屋分のvisited配列をFalseで初期化します。
• queue := addRooms(rooms, 0, queue, visited) を実行します。
• visited[0] := True と設定します。
• キューに要素が存在する間、以下を繰り返します。
  • queue := addRooms(rooms, queue[0], queue, visited)
  • visited[queue[0]] を True にマークします。
  • キューの先頭要素を削除します。
• 最後に、visited配列のすべての要素がTrueであればtrueを返します。

add_rooms()関数の処理内容

• 引数としてrooms、index、queue、visited配列を受け取ります。
• rooms[index]内の各要素iについて、
  • iが未訪問であれば、キューに挿入します。
• キューを返します。

Pythonでの実装例

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

class Solution(object):
    def canVisitAllRooms(self, rooms):
        queue = []
        visited = [False for i in rooms]
        queue = self.add_rooms(rooms,0,queue,visited)
        visited[0] = True
        while len(queue)>0:
            queue = self.add_rooms(rooms,queue[0],queue,visited)
            visited[queue[0]] = True
            queue.pop(0)
        return all(visited)
    def add_rooms(self, rooms,index,queue,visited):
        for i in rooms[index]:
            if not visited[i]:
                queue.append(i)
        return queue
ob1 = Solution()
print(ob1.canVisitAllRooms([[1],[2],[3],[]]))

入力

[[1],[2],[3],[]]

出力

true

このコードでは、キューを用いて訪問予定の部屋を管理しながら、各部屋から入手した鍵を順次処理していきます。最終的にall(visited)によって全部屋を訪問できたかどうかを一括判定しており、計算量は部屋数と鍵の総数に比例するO(N + K)で効率的です。