[ BOJ ] 17472 : 다리 만들기 2 ( GOLD 1 ) / Python

>> 문제 바로가기 (https://www.acmicpc.net/problem/17472)

문제

섬으로 이루어진 나라가 있고, 모든 섬을 다리로 연결하려고 한다. 이 나라의 지도는 N×M 크기의 이차원 격자로 나타낼 수 있고, 격자의 각 칸은 땅이거나 바다이다.

섬은 연결된 땅이 상하좌우로 붙어있는 덩어리를 말하고, 아래 그림은 네 개의 섬으로 이루어진 나라이다. 색칠되어있는 칸은 땅이다.

다리는 바다에만 건설할 수 있고, 다리의 길이는 다리가 격자에서 차지하는 칸의 수이다. 다리를 연결해서 모든 섬을 연결하려고 한다. 섬 A에서 다리를 통해 섬 B로 갈 수 있을 때, 섬 A와 B를 연결되었다고 한다. 다리의 양 끝은 섬과 인접한 바다 위에 있어야 하고, 한 다리의 방향이 중간에 바뀌면 안된다. 또, 다리의 길이는 2 이상이어야 한다.

다리의 방향이 중간에 바뀌면 안되기 때문에, 다리의 방향은 가로 또는 세로가 될 수 밖에 없다. 방향이 가로인 다리는 다리의 양 끝이 가로 방향으로 섬과 인접해야 하고, 방향이 세로인 다리는 다리의 양 끝이 세로 방향으로 섬과 인접해야 한다.

섬 A와 B를 연결하는 다리가 중간에 섬 C와 인접한 바다를 지나가는 경우에 섬 C는 A, B와 연결되어있는 것이 아니다. 

아래 그림은 섬을 모두 연결하는 올바른 2가지 방법이고, 다리는 회색으로 색칠되어 있다. 섬은 정수, 다리는 알파벳 대문자로 구분했다.

다음은 올바르지 않은 3가지 방법이다.

다리가 교차하는 경우가 있을 수도 있다. 교차하는 다리의 길이를 계산할 때는 각 칸이 각 다리의 길이에 모두 포함되어야 한다. 아래는 다리가 교차하는 경우와 기타 다른 경우에 대한 2가지 예시이다.

나라의 정보가 주어졌을 때, 모든 섬을 연결하는 다리 길이의 최솟값을 구해보자.

입력

첫째 줄에 지도의 세로 크기 N과 가로 크기 M이 주어진다. 둘째 줄부터 N개의 줄에 지도의 정보가 주어진다. 각 줄은 M개의 수로 이루어져 있으며, 수는 0 또는 1이다. 0은 바다, 1은 땅을 의미한다.

출력

모든 섬을 연결하는 다리 길이의 최솟값을 출력한다. 모든 섬을 연결하는 것이 불가능하면 -1을 출력한다.

풀이 과정

다음 과정을 거쳐 구현했다.

 

1. BFS 로 섬 영역을 탐색하고, 영역마다 번호를 부여

2. 섬부터 섬까지 다리를 가로 세로로 지었을 때의 거리를 측정하고, 다리 길이가 2 이상이라면 간선으로 등록

3. Kruskal MST 로 최소 신장 트리 완성

4. 모든 노드가 연결되었는지 유니온 파인드로 확인

 

1. BFS 로 섬 영역을 탐색하고, 영역마다 번호를 부여

visited_mapnum = [[0 for _ in range(m)] for _ in range(n)]
row = [-1, 1, 0, 0]
col = [0, 0, -1, 1]

node_num = 1
node_map = [[0 for _ in range(m)] for _ in range(n)]
for i in range(n):
    for j in range(m):
        if maps[i][j] == 1 and visited_mapnum[i][j] == 0:
            visited_mapnum[i][j] = 1
            node_map[i][j] = node_num
            queue = deque()
            queue.append([i, j])

            while queue:
                now = queue.popleft()

                for d in range(4):
                    ny = now[0] + row[d]
                    nx = now[1] + col[d]

                    if 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 1 and visited_mapnum[ny][nx] == 0:
                        visited_mapnum[ny][nx] = 1
                        node_map[ny][nx] = node_num
                        queue.append([ny, nx])

            node_num += 1
node_num -= 1

BFS 를 통해 섬 영역을 탐색하고, 번호를 부여하여 노드처럼 다룰 준비를 했다.

 

2. 섬부터 섬까지 다리를 가로 세로로 지었을 때의 거리를 측정

graph = [{} for _ in range(node_num + 1)]
for i in range(n):
    for j in range(m):
        if maps[i][j] == 1:
            for d in range(4):
                ny = i + row[d]
                nx = j + col[d]

                if 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 0:
                    cnt = 1
                    while 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 0:
                        ny += row[d]
                        nx += col[d]
                        cnt += 1
                    
                    if 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 1 and node_map[i][j] != node_map[ny][nx] and cnt - 1 >= 2:
                        if graph[node_map[i][j]].get(node_map[ny][nx], -1) == -1:    
                            graph[node_map[i][j]][node_map[ny][nx]] = cnt - 1
                        else:
                            graph[node_map[i][j]][node_map[ny][nx]] = min(graph[node_map[i][j]][node_map[ny][nx]], cnt - 1)
                        
                        if graph[node_map[ny][nx]].get(node_map[i][j], -1) == -1:
                            graph[node_map[ny][nx]][node_map[i][j]] = cnt - 1
                        else:
                            graph[node_map[ny][nx]][node_map[i][j]] = min(graph[node_map[ny][nx]][node_map[i][j]], cnt - 1)

특정 섬 지점부터 시작해서 다리를 상하좌우로 쭉 지었을 때, ( + 다른 섬을 지나지 않을 때까지 ) 다리 길이를 카운트한다.

같은 섬끼리 이어질 수도 있으니 주의해야 한다는 점, A 섬 -> B 섬 다리가 여러 개일 수도 있으니 최솟값 처리를 해줘야 한다는 점을 유의한다.

 

3. Kruskal MST 로 최소 신장 트리 완성

arr = [i for i in range(node_num + 1)]
def parent(now):
    if now != arr[now]:
        arr[now] = parent(arr[now])
    return arr[now]

def union(a, b):
    a = parent(a)
    b = parent(b)

    if a < b: arr[b] = a
    else: arr[a] = b

edges = []
for a in range(1, node_num + 1):
    for b, cost in list(graph[a].items()):
        edges.append([cost, a, b])

edges.sort()
result = 0

for i in range(len(edges)):
    nc, na, nb = edges[i]
    if parent(na) != parent(nb):
        union(na, nb)
        result += nc

모든 섬을 연결할 수 있는 최소 신장 트리를 만든다.

 

4. 모든 노드가 연결되었는지 유니온 파인드로 확인

for i in range(1, node_num + 1): arr[i] = parent(i)
if arr.count(arr[1]) == node_num: print(result)
else: print(-1)

특정 섬으로 갈 수 있는 다리가 없다는 등, 모든 섬이 연결이 되지 않을 수도 있으므로 마지막으로 전부 부모 관계를 갱신해준 뒤, 모든 노드가 동일한 부모를 가지고 있다면 MST 결과를 출력한다.

 

전체 코드

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

n, m = map(int, input().rstrip().split())
maps = [list(map(int, input().rstrip().split())) for _ in range(n)]

# 1. BFS 로 영역 탐색 및 노드 번호 부여
# 2. A -> B 영역 거리 탐색 및 기록 ( 2 이상 거리, 가로 세로 여야 함 )
# 3. kruskal MST 로 모든 노드가 유파로 연결 되었는지 체크

visited_mapnum = [[0 for _ in range(m)] for _ in range(n)]
row = [-1, 1, 0, 0]
col = [0, 0, -1, 1]

node_num = 1
node_map = [[0 for _ in range(m)] for _ in range(n)]
for i in range(n):
    for j in range(m):
        if maps[i][j] == 1 and visited_mapnum[i][j] == 0:
            visited_mapnum[i][j] = 1
            node_map[i][j] = node_num
            queue = deque()
            queue.append([i, j])

            while queue:
                now = queue.popleft()

                for d in range(4):
                    ny = now[0] + row[d]
                    nx = now[1] + col[d]

                    if 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 1 and visited_mapnum[ny][nx] == 0:
                        visited_mapnum[ny][nx] = 1
                        node_map[ny][nx] = node_num
                        queue.append([ny, nx])

            node_num += 1
node_num -= 1

graph = [{} for _ in range(node_num + 1)]
for i in range(n):
    for j in range(m):
        if maps[i][j] == 1:
            for d in range(4):
                ny = i + row[d]
                nx = j + col[d]

                if 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 0:
                    cnt = 1
                    while 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 0:
                        ny += row[d]
                        nx += col[d]
                        cnt += 1
                    
                    if 0 <= ny < n and 0 <= nx < m and maps[ny][nx] == 1 and node_map[i][j] != node_map[ny][nx] and cnt - 1 >= 2:
                        if graph[node_map[i][j]].get(node_map[ny][nx], -1) == -1:    
                            graph[node_map[i][j]][node_map[ny][nx]] = cnt - 1
                        else:
                            graph[node_map[i][j]][node_map[ny][nx]] = min(graph[node_map[i][j]][node_map[ny][nx]], cnt - 1)
                        
                        if graph[node_map[ny][nx]].get(node_map[i][j], -1) == -1:
                            graph[node_map[ny][nx]][node_map[i][j]] = cnt - 1
                        else:
                            graph[node_map[ny][nx]][node_map[i][j]] = min(graph[node_map[ny][nx]][node_map[i][j]], cnt - 1)

arr = [i for i in range(node_num + 1)]
def parent(now):
    if now != arr[now]:
        arr[now] = parent(arr[now])
    return arr[now]

def union(a, b):
    a = parent(a)
    b = parent(b)

    if a < b: arr[b] = a
    else: arr[a] = b

edges = []
for a in range(1, node_num + 1):
    for b, cost in list(graph[a].items()):
        edges.append([cost, a, b])

edges.sort()
result = 0

for i in range(len(edges)):
    nc, na, nb = edges[i]
    if parent(na) != parent(nb):
        union(na, nb)
        result += nc

for i in range(1, node_num + 1): arr[i] = parent(i)
if arr.count(arr[1]) == node_num: print(result)
else: print(-1)