[백준] 17472 - 다리 만들기 2

https://www.acmicpc.net/problem/17472

구현 + MST

 

섬의 위치가 격자판 모양으로 주어질 때, 격자판 모양으로부터 그래프를 직접 정의하고, 그 그래프에 대해 MST를 찾으면 됩니다.

이런 복잡한 문제는 어떤 것들을 구현해야 하는지 잘 나눠서 생각하면 구현하기 좋습니다.

(그래도 이 문제는 구현이 엄청 복잡한 편까지는 아닙니다.)

 

먼저 섬부터 찾아보겠습니다.

섬을 찾는다는 것은 격자판을 그대로 그래프로 보았을 때, 컴포넌트를 찾는 것과 같습니다.

# 섬 찾기
island_number = 2
island_tiles = []
for i in range(N):
    for j in range(M):
        if board[i][j] == 1:
            # BFS
            d = deque()
            d.append((i, j))
            board[i][j] = island_number
            island_tiles.append((i, j))

            while d:
                now_i, now_j = d.popleft()
                for k in range(4):
                    check_i = now_i + di[k]
                    check_j = now_j + dj[k]
                    if 0 <= check_i < N and 0 <= check_j < M:
                        if board[check_i][check_j] == 1:
                            board[check_i][check_j] = island_number
                            d.append((check_i, check_j))
                            island_tiles.append((check_i, check_j))

            island_number += 1

 

그래프에서 BFS를 돌면서 1을 만날 때마다 그 1에 대해서 BFS를 돌았습니다.

저는 BFS 구현이 익숙해서 BFS로 돌았는데, DFS로 돌아도 상관없습니다.

 

각각의 섬마다 번호를 붙이고 새로운 그래프를 구성해야 하기 때문에, 섬에 2부터 번호를 붙여주었습니다.

번호를 1부터 붙이지 않은 이유는 이미 격자판에서 1이라는 숫자를 사용하고 있었기 때문입니다.

저는 1이라는 숫자에 아직 방문하지 않은 섬이라는 의미까지 포함시켜서 방문 체크용으로도 사용했습니다.

(사실 섬의 번호, 방문 여부 용도가 한 숫자에 섞여서 쓰이는 것은 나중에 디버깅을 어렵게 하는 요소로 작용해서 좋은 구현은 아니라고 생각합니다.)

 

# edge 찾기
def find_island(now_i, now_j, now_dir):
    count = 0
    while 0 <= now_i < N and 0 <= now_j < M and board[now_i][now_j] == 0:
        count += 1
        now_i += di[now_dir]
        now_j += dj[now_dir]

    if 0 <= now_i < N and 0 <= now_j < M and count > 1:
        return board[now_i][now_j], count
    return -1, -1

edge = []
for i, j in island_tiles:
    now_island = board[i][j]
    for k in range(4):
        check_i = i + di[k]
        check_j = j + dj[k]
        if 0 <= check_i < N and 0 <= check_j < M and board[check_i][check_j] == 0:
            another_island, dist = find_island(check_i, check_j, k)
            if another_island != -1:
                edge.append((dist, now_island, another_island))

 

다음으로 엣지를 찾습니다.

보드의 최대 크기가 작아서, 다시 보드 전체를 돌아도 되지만, 섬을 찾는 과정에서 이미 보드를 돌았기 때문에, 이 과정에서 섬이 되는 노드들을 island_tiles 에 저장해 두도록 했습니다.

 

저장한 island_tiles를 돌면서, 해당 섬 타일 주변에 바다(0)가 있는지 확인합니다.

바다가 있다면 해당 방향을 향해 바다가 아닌 타일이 나올 때까지 직진합니다.

바다가 아닌 타일을 만나면, 그때까지 만난 바다 타일의 개수를 거리로 하는 edge를 반환합니다.

문제 조건에 따라 만약 바다 타일의 개수가 1개 이하라면 그때는 edge가 없는 것으로 봅니다.

 

이렇게 발견한 엣지는 같은 섬 쌍에 대해서도 모두 edge로 추가했습니다.

보드의 크기가 최대 10 * 10 으로 작아서 이렇게 저장해도 edge의 전체 개수가 많지 않기 때문입니다.

 

만약 같은 섬 쌍에 대해서 하나의 엣지만 저장하고 싶다면, 꼭 두 섬 사이를 잇는 가장 짧은 엣지만 저장하도록 구현해야 합니다.

---

# MST
node = island_number
edge.sort()
parent = [i for i in range(node)]


def find(x):
    if parent[x] == x:
        return x
    v = find(parent[x])
    parent[x] = v
    return v


def union(x, y):
    x = find(x)
    y = find(y)
    parent[y] = x


answer = 0
conn_count = 0
for dist, island1, island2 in edge:
    island1 = find(island1)
    island2 = find(island2)
    if island1 != island2:
        answer += dist
        union(island1, island2)
        conn_count += 1

island_count = node - 2
if conn_count == island_count - 1:
    print(answer)
else:
    print(-1)

 

이제 그래프를 만들었으니, 새로 정의한 그래프에 대해 MST를 돌기만 하면 됩니다.

저는 크루스칼 알고리즘을 이용해서 돌았습니다.

 

이때 또 조심해야 할 점은 '모든 섬을 연결하는' MST를 찾아야 하기 때문에, 몇 개의 섬을 연결했는지 체크해야 합니다.

크루스칼 알고리즘은 n개 노드로 구성된 MST를 만들 때, n-1 번의 union 연산을 호출하는 점을 이용해서 몇 개 노드로 구성된 MST인지 체크하였습니다.

 

만약 MST의 노드 개수가 섬의 개수와 같다면 전체 간선의 합을 출력하고, 아니라면 -1을 출력하도록 하였습니다.

 

---

import sys
from collections import deque
input = sys.stdin.readline

di = [0, 0, -1, 1]
dj = [-1, 1, 0, 0]

N, M = map(int, input().split())
board = [list(map(int, input().split())) for _ in range(N)]
visit = [[False for _ in range(M)] for _ in range(N)]

# 섬 찾기
island_number = 2
island_tiles = []
for i in range(N):
    for j in range(M):
        if board[i][j] == 1:
            # BFS
            d = deque()
            d.append((i, j))
            board[i][j] = island_number
            island_tiles.append((i, j))

            while d:
                now_i, now_j = d.popleft()
                for k in range(4):
                    check_i = now_i + di[k]
                    check_j = now_j + dj[k]
                    if 0 <= check_i < N and 0 <= check_j < M:
                        if board[check_i][check_j] == 1:
                            board[check_i][check_j] = island_number
                            d.append((check_i, check_j))
                            island_tiles.append((check_i, check_j))

            island_number += 1

# edge 찾기
def find_island(now_i, now_j, now_dir):
    count = 0
    while 0 <= now_i < N and 0 <= now_j < M and board[now_i][now_j] == 0:
        count += 1
        now_i += di[now_dir]
        now_j += dj[now_dir]

    if 0 <= now_i < N and 0 <= now_j < M and count > 1:
        return board[now_i][now_j], count
    return -1, -1

edge = []
for i, j in island_tiles:
    now_island = board[i][j]
    for k in range(4):
        check_i = i + di[k]
        check_j = j + dj[k]
        if 0 <= check_i < N and 0 <= check_j < M and board[check_i][check_j] == 0:
            another_island, dist = find_island(check_i, check_j, k)
            if another_island != -1:
                edge.append((dist, now_island, another_island))


# MST
node = island_number
edge.sort()
parent = [i for i in range(node)]


def find(x):
    if parent[x] == x:
        return x
    v = find(parent[x])
    parent[x] = v
    return v


def union(x, y):
    x = find(x)
    y = find(y)
    parent[y] = x


answer = 0
conn_count = 0
for dist, island1, island2 in edge:
    island1 = find(island1)
    island2 = find(island2)
    if island1 != island2:
        answer += dist
        union(island1, island2)
        conn_count += 1

island_count = node - 2
if conn_count == island_count - 1:
    print(answer)
else:
    print(-1)

 

전체 코드는 다음과 같습니다.