【Python】森のすべての木が燃え尽きるまでの日数を求めるアルゴリズム

問題の概要

2次元の行列で森を表すことを考えます。各マスは次の3種類のいずれかです。

• 0:空き地(何もないマス)
• 1:木のあるマス
• 2:燃えている木のマス

毎日、上下左右に隣接するマス(斜め方向は含まない)の木が燃えていると、その木にも火が燃え移ります。このときすべての木が燃え尽きるまでにかかる日数を求めてください。もし全部の木を燃やすことが不可能な場合は -1 を返します。

入力例

たとえば、次のような森が与えられたとします。

この場合の出力は 4 になります。上段中央の燃えている木から火が1日ごとに周囲へ広がり、4日目についに最後の木が燃え尽きるためです。

解き方の考え方

この問題は「複数起点の幅優先探索(BFS)」と呼ばれる典型的なパターンです。有名な「腐ったオレンジ」の問題と同じ発想で、最初に燃えている木をすべてキューに入れ、1日ごとに同時に火を広げていくことで最小日数を求められます。

具体的な手順は以下の通りです。

• ans := 0(経過日数のカウンタ)
• twos := 新しいリスト(現在燃えているマスの座標を格納)
• 行列全体を走査し、値が 2 のマスの座標 (i, j) をすべて twos に追加する
• twos が空でない間、以下を繰り返す:
  • temp := 新しいリスト(次の日に新しく燃えるマス)
  • twos 内の各座標 (i, j) について、4近傍 [(i+1, j), (i, j+1), (i−1, j), (i, j−1)] のそれぞれ (x, y) に対して:
    • (x, y) が行列の範囲内であり、かつ matrix[x][y] が 1 なら、(x, y) を temp の末尾に追加する
  • temp 内の各座標について matrix[i][j] := 2 と更新する(火が燃え移る)
  • twos := temp と置き換える
  • twos が空でなければ ans := ans + 1(1日経過)
• 最後に、行列内に残っている 1(未燃の木)の個数 ones を数える
• ones が 0 なら ans を返し、そうでなければ -1 を返す

最後に残りの木を数えているのは重要なポイントです。空き地(0)に阻まれて火が届かない木は永遠に燃えないため、その場合は -1 を返す必要があります。

実装例(Python)

それでは、実際のコードを見てみましょう。

class Solution:
    def solve(self, matrix):
        ans = 0
        twos = []
        for i in range(len(matrix)):
            for j in range(len(matrix[0])):
                if matrix[i][j] == 2:
                    twos.append((i, j))
        while twos:
            temp = []
            for i, j in twos:
                for x, y in [(i + 1, j), (i, j + 1), (i - 1, j), (i, j - 1)]:
                    if 0 <= x < len(matrix) and 0 <= y < len(matrix[0]) and matrix[x][y] == 1:
                        temp.append((x, y))
            for i, j in temp:
                matrix[i][j] = 2
            twos = temp
            ans += 1 if twos else 0

        ones = sum(int(matrix[i][j] == 1) for i in range(len(matrix)) for j in range(len(matrix[0])))
        return ans if ones == 0 else -1

ob = Solution()
matrix = [
    [1, 2, 1],
    [1, 0, 1],
    [1, 1, 1]
]
print(ob.solve(matrix))

入力

matrix = [
[1, 2, 1],
[1, 0, 1],
[1, 1, 1]
]

出力

4

計算量について

各行・各マスは最大でも1回しかキューに追加されないため、時間計算量は O(R × C)(Rは行数、Cは列数)となり、非常に効率的です。空間計算量も同様に O(R × C) です。