01 Matrix (Leetcode)

今回の問題

leetcodeの01 Matrixをやりました。問題はhttps://leetcode.com/problems/01-matrix/
を参照してください。
ざっくり説明すると、0か1で埋まった二次元配列が渡されるのでそれぞれのますについて0との最短距離を求めるという問題です。

元の方針

最初に立てた方針としては

1. 与えられた二次元配列と同じ大きさの配列を作る（0で初期化)
2. 配列を一つずつ見ていって1ならば別で作った再帰関数に入る。
3. 再帰関数では上下左右のマスを見る。0なら1を返す。1ならまた再帰。最初に入った上下左右の返り値で最小のものを配列に代入。

コードは以下のようにしました。

class Solution:
    def nearest(self, mat, i, j):
        if i >= 0 and i < len(mat) and j >= 0 and j < len(mat[0]):
            if mat[i][j] == 0:
                return 0
            else:
                return 1 + min(self.nearest(mat, i-1, j), self.nearest(mat, i, j-1), self.nearest(mat, i+1, j), self.nearest(mat, i, j+1))
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        result = [[0] * len(mat[0]) for i in range(len(mat))]
        for i in range(len(mat)):
            for j in range(len(mat[0])):
                if mat[i][j] == 1:
                    result[i][j] = self.nearest(mat, i, j)
        return result

ただこれはうまくいかなかった。再帰関数で上下左右見る時に1が一つでもあると無限ループに入ってしまう。

幅優先探索

なので、幅優先探索を用いて実装することにしました。

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        col_len = len(mat[0])
        row_len = len(mat)
        move = [[-1, 0], [0, 1], [1, 0], [0, -1]]
        queue = []
        for i in range(row_len):
            for j in range(col_len):
                if mat[i][j] == 0:
                    queue.append([i, j])
                else:
                    mat[i][j] = "1"
        for i, j in queue:
            for x, y in move:
                col = j + y
                row = i + x
                if col >= 0 and col < col_len and row >= 0 and row < row_len:
                    if mat[row][col] == "1":
                        mat[row][col] = mat[i][j] + 1
                        queue.append([row, col])
        return mat