Logo

01 Matrix

LeetCode의 542번째 문제인 01 Matrix를 함께 풀어보도록 하겠습니다.

문제

주어진 m x n 이차원 배열 mat에서 각 칸과 가장 가까운 0 간의 거리를 반환하시오. 인접한 두 칸 사이의 거리는 1입니다.

예제

Input: mat = [[0,0,0],[0,1,0],[0,0,0]]
Output: [[0,0,0],[0,1,0],[0,0,0]]
Input: mat = [[0,0,0],[0,1,0],[1,1,1]]
Output: [[0,0,0],[0,1,0],[1,2,1]]

풀이 1

주어진 행렬에서 간 칸과 가장 가까운 0을 어떻게 찾을 수 있을까요?

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

첫 번째 예제를 가지고 생각해보면, 행렬의 정 중앙에 있는 1은 8개의 0으로 둘러싸여 있는데요. 상하좌우에 있는 0까지는 한 칸만 이동하면 되니까 거리가 1이고요. 네 모서리에 있는 0까지는 두 칸을 이동해야하므로 거리가 2가 됩니다. 따라서, 정 중앙에 있는 (1, 1) 위치에서 가장 가까운 0과의 거리는 1이 되죠.

문제에서 주어진 행렬은 숫자가 들어있는 각 칸을 정점(Vertex, Node), 상하좌우에 있는 칸과의 연결을 간선(Edge)로 보면 그래프(Graph) 자료구조로 나타낼 수 있는데요. 있으니까요. 값이 0인 모든 정점에서부터 그래프 탐색을 시작해서 값이 0이 아닌 정점에 도달할 때마다 두 지점 간의 거리를 잴 수 있을 것입니다. 그런데 우리는 이 중에서 최소 거리를 구해야 하기 때문에 이전에 해당 정점에서 구했던 다른 거리와 비교하여 가장 짧은 거리를 선택해야 하겠죠?

그럼 다음과 같이 0이 세 개 들어있는 행렬을 이용해서 알고리즘을 좀 더 구체화해보겠습니다.

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

추가 행렬을 만들기 보다는 입력 행렬에 최소 거리를 바로 저장하는 것이 간단할 것 같습니다. 거리를 구할 때 마다 더 작은 값으로 갱신이 용이하도록, 입력 행렬에서 0이 아닌 칸의 값을 가능한 최대 거리로 초기화해 주겠습니다. 3개의 행과 3개의 열로 이루어진 행렬에서 칸 사이의 거리는 아무리 크더라도 절대 3 x 3를 넘지 못할 것이므로 9로 갱신해줄께요.

[0, 0, 9]
[9, 0, 9]
[9, 9, 9]

첫 번째 0이 있는 위치인 (0, 0)부터 거리를 재보면 다음과 같은데요.

[0, 1, 2]
[1, 2, 3]
[2, 3, 4]

이 거리를 각 칸에 있던 거리와 비교를 해서 더 작은 거리로 갱신해줍니다.

[0, 0, 2]
[1, 0, 3]
[2, 3, 4]

두 번째 0이 있는 위치인 (0, 1)부터 거리를 재보면 다음과 같은데요.

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

이 거리를 각 칸에 있던 거리와 비교를 해서 더 작은 거리로 갱신해줍니다.

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

세 번째 0이 있는 위치인 (1, 1)부터 거리를 재보면 다음과 같은데요.

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

이 거리를 각 칸에 있던 거리와 비교를 해서 더 작은 거리로 갱신해줍니다.

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

결국은, 3개의 0부터 잰 거리 중에서 최소 값을 선택한 효과가 납니다.

[min(0, 1, 2), min(1, 0, 1), min(2, 1, 2)]
[min(1, 2, 1), min(2, 1, 0), min(2, 1, 2)]
[min(2, 3, 2), min(3, 2, 1), min(4, 3, 2)]

이 깊이 우선 탐색을 재귀 알고리즘을 이용하여 파이썬으로 구현해보겠습니다.

class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        n_rows, n_cols = len(mat), len(mat[0])

        def dfs(row, col):
            for r, c in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                if 0 <= row + r < n_rows and 0 <= col + c < n_cols:
                    if mat[row + r][col + c] > mat[row][col] + 1:
                        mat[row + r][col + c] = mat[row][col] + 1
                        dfs(row + r, col + c)

        for r in range(n_rows):
            for c in range(n_cols):
                if mat[r][c] != 0:
                    mat[r][c] = n_rows * n_cols

        for r in range(n_rows):
            for c in range(n_cols):
                if mat[r][c] == 0:
                    dfs(r, c)

        return mat

주어진 입력 행렬의 행의 수를 r, 열의 수를 c라고 했을 때, 이 풀이의 시간 복잡도는 O((r * c)^2)이 되는데요. (2, 0)(2, 2) 경우처럼, 하나의 칸을 두 번 이상 방문하여 더 작은 거리로 갱신할 수 있기 때문입니다. 공간 복잡도는 재귀 함수의 호출 스택의 최대 깊이를 생각해보면 O(r * c)라는 것을 알 수 있습니다.

풀이 2

깊이 우선 탐색을 하게 되면 0인 칸에서 0이 아닌 칸까지 도달할 수 있는 모든 경우의 수를 따져봐야하는데요. 이 문제는 결국 그래프 내에서 정점 간의 최소 거리를 구하는 문제이므로 너비 우선 탐색을 하면 훨씬 효과적으로 해결할 수 있을 것 같습니다.

기본 아이디어는 주어진 행렬 내에서 0이 들어있는 칸부터 시작해서 0이 들어있지 않는 칸을 향해서 거리를 1씩 점점 늘려가면서 탐색을 하는 건데요. 이렇게 하면 0이 아닌 셀에 제일 처음 도달했을 때의 거리가 무조건 0부터 해당 셀까지의 가장 가까운 거리라고 확신할 수 있습니다. 그러므로 0이 들어있는 다른 셀에서 해당 셀까지 도착할 수 있는 경로를 모두 배재해버릴 수 있겠죠? 즉, 너비 우선 탐색을 통해서 연산량을 획기적으로 줄이고 훨씬 빨리 최소 거리를 찾을 수 있습니다.

그럼 문제에서 주어진 두 번째 예제로 다시 생각을 해볼까요? 우선 0이 들어있는 모든 칸의 좌표를 큐(Queue) 자료구조에 추가합니다.

큐: [(0, 0), (0, 1), (0, 2), (1, 0), (1, 2)]

큐에서 (0, 0)을 제거하고, 행렬에서 이 위치 기준으로 갱신할 칸이 있는지 사방을 살핍니다. 위와 왼쪽은 칸이 없고, 아래와 오른쪽 칸에는 이미 최소 거리인 0이 들어가 있습니다.

(0, 0) 제거

큐: [(0, 1), (0, 2), (1, 0), (1, 2)]
행렬:
[0, 0, 0]
[0, ?, 0]
[?, ?, ?]

다음에는 큐에서 (0, 1)을 제거하고, 행렬에서 이 위치 기준으로 사방을 살핍니다. 위에는 칸이 없고, 좌우 칸에는 이미 최소 거리 0이 들어있으며, 아래 칸에는 ?가 들어있는데요. 따라서 여기서 부터 아래 칸과의 거리인 1로 갱신해줍니다. 나중에 아래 칸에서 부터 거리를 젤 수 있도록 아래 칸의 위치인 (1, 1)을 큐에 추가합니다.

(0, 1) 제거
mat[1][1] = mat[0][1] + 1 = 0 + 1 = 1
(1, 1) 추가

큐: [(0, 2), (1, 0), (1, 2), (1, 1)]
행렬:
[0, 0, 0]
[0, 1, 0]
[?, ?, ?]

다음에는 큐에서 (0, 2)를 제거하고, 행렬에서 이 위치 기준으로 사방을 살핍니다. 위와 오른쪽에는 칸이 없고, 아래와 왼쪽에는 0이 들어있네요. 그럼 굳이 여기서 부터 거리를 젤 이유가 없으므로 넘어갑니다.

(0, 2) 제거

큐: [(1, 0), (1, 2), (1, 1)]
행렬:
[0, 0, 0]
[0, 1, 0]
[?, ?, ?]

다음에는 큐에서 (1, 0)을 제거하고, 행렬에서 이 위치 기준으로 사방을 살핍니다. 왼쪽에는 칸이 없고, 위에 있는 칸에는 0이 있고 우측 칸에는 1이 들어있는데, 모두 이미 최소 거리일 것입니다. ?가 들어있는 아래 칸은 여기서부터 거리를 재면 1로 갱신할 수 있습니다. 나중에 아래 칸에서 부터 거리를 젤 수 있도록 아래 칸의 위치인 (2, 0)을 큐에 추가합니다.

(1, 0) 제거
mat[2][0] = mat[1][0] + 1 = 0 + 1 = 1
(2, 0) 추가

큐: [(1, 2), (1, 1), (2, 0)]
행렬:
[0, 0, 0]
[0, 1, 0]
[1, ?, ?]

큐에서 (1, 2)을 제거하고, 이 위치 기준으로 사방을 살핍니다. 아래 칸에 ?가 들어있으니, 여기서 부터 잰 거리인 1로 갱신할 수 있습니다.

(1, 2) 제거
mat[2][2] = mat[1][2] + 1 = 0 + 1 = 1
(2, 2) 추가

큐: [(1, 1), (2, 0), (2, 2)]
행렬:
[0, 0, 0]
[0, 1, 0]
[1, ?, 1]

큐에서 (1, 1)을 제거하고, 이 위치 기준으로 사방을 살핍니다. 상, 좌, 우는 이미 최소 거리인 0이지만, 아래에는 ?가 있습니다. 아래 칸을 여기서부터 거리를 재면 2로 갱신할 수 있습니다. 나중에 아래 칸에서 부터 거리를 젤 수 있도록 아래 칸의 위치인 (2, 1)을 큐에 추가합니다.

(1, 1) 제거
mat[2][1] = mat[1][1] + 1 = 1 + 1 = 2
(2, 1) 추가

큐: [(2, 0), (2, 2), (2, 1)]
행렬:
[0, 0, 0]
[0, 1, 0]
[1, 2, 1]

큐에서 (2, 0)을 제거하고, 사방을 살피면 모두 이미 최소 거리를 담고 있습니다.

(2, 0) 제거

큐: [(2, 2), (2, 1)]
행렬:
[0, 0, 0]
[0, 1, 0]
[1, 2, 1]

큐에서 (2, 2)을 제거하고, 사방을 살피면 모두 이미 최소 거리를 담고 있습니다.

(2, 2) 제거

[(2, 1)]
행렬:
[0, 0, 0]
[0, 1, 0]
[1, 2, 1]

큐에서 (2, 1)을 제거하고, 이 위치 기준으로 사방을 살핍니다. 상, 좌, 우 모두 1이 있으며, 모두 이미 최소 거리를 담고 있으므로 넘어갑니다.

(2, 1) 제거

[]
행렬:
[0, 0, 0]
[0, 1, 0]
[1, 2, 1]

이제 큐가 비었기 때문에 더 이상 거리를 따져볼 좌표가 없습니다.

지금까지 설명드린 너비 우선 탐색 알고리즘을 파이썬으로 구현해볼까요? 큐(Queue) 자료구조가 필요하기 때문에 파이썬에 내장된 collections 모듈의 deque를 사용하겠습니다.

from collections import deque


class Solution:
    def updateMatrix(self, mat: List[List[int]]) -> List[List[int]]:
        n_rows, n_cols = len(mat), len(mat[0])
        queue = deque()

        for r in range(n_rows):
            for c in range(n_cols):
                if mat[r][c] == 0:
                    queue.append((r, c))
                else:
                    mat[r][c] = -1 # 아직 구하지 않은 최소 거리

        while queue:
            row, col = queue.popleft()

            for r, c in [(-1, 0), (1, 0), (0, -1), (0, 1)]:
                if 0 <= row + r < n_rows and 0 <= col + c < n_cols:
                    if mat[row + r][col + c] == -1:
                        mat[row + r][col + c] = mat[row][col] + 1
                        queue.append((row + r, col + c))

        return mat

이 풀이의 시간 복잡도는 O(r * c)로 향상이 됩니다. 왜냐하면 각 좌표는 큐에 딱 한 번만 들어갔다가 나오기 때문입니다. 공간 복잡도는 큐에 최대 r * c 개의 좌표가 저장될 수 있으므로 O(r * c)이 됩니다.

마치면서

코딩 시험에서 그래프(Graph)을 다루는 유형의 문제에서는 이 문제가 가장 기본이 된다고 볼 수 있는데요. 이 문제가 너무 쉬우셨다면 비슷하지만 좀 더 어려운 문제인 Flood Fill도 풀어보시라고 추천드립니다. 코딩 테스트에서 그래프를 어떻게 다루는지에 대해서 더 공부하고 싶으시다면 관련 게시물을 참고 바랄께요.

본 문제를 통해 기본기를 잘 닦아놓으셔서 같은 유형의 좀 더 어려운 문제를 풀 때 큰 도움이 되었으면 좋겠습니다.