[Python 3] BOJ - 1944 "복제 로봇"

 # 문제 링크

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

1944번: 복제 로봇

 

 # 풀이

 개인적으로 이 문제의 핵심은 최소 스패닝 트리과 그래프 탐색이라고 생각한다. 

<그림 1. 예제 입력 1의 경우>

 위 그림 1의 좌측 그림은 로봇이 자기자신을 복사할 수 없는 상황에서 모든 열쇠를 찾는 경우이고, 오른쪽 그림은 로봇이 문제의 조건에 따라 자기자신을 복사할 수 있는 상황에서 모든 열쇠를 찾는 경우이다. 왼쪽의 경우 빨간색, 파란색, 초록색, 노란색 화살표 순서대로 방문해야 모든 열쇠를 얻을 수 있으나, 오른쪽의 경우 현재 진행하는 방향과 다른 방향으로 갈 수 있다면 복제가 가능하므로 굳이 돌아올 필요가 없다. 즉, 오른쪽의 경우는 왼쪽의 파란색 화살표가 필요 없다는 것이다.

 그런데 문제에서는 모든 로봇들이 움직이는 거리의 합을 최소화 시켜야 하고 모든 열쇠를 찾아야 하므로 현재 입력으로 주어진 미로에서 모든 S와 K는 연결되어야 한다. 모든 정점 쌍에 대해 경로가 존재하며, 간선의 합이 최소화하는 것은 정확히 최소 스패닝 트리의 정의와 동치이므로 입력으로 주어진 2차원 배열에서 S와 K를 정점으로 하는 그래프로 모델링 후 최소 스패닝 트리를 구할 수 있다면 그때의 값을, 구할 수 없다면 -1을 출력하면 정답을 구할 수 있다. 다만, 이 때 어떤 정점에서 다른 정점으로 가는 경우가 여러개 있으므로 모든 S와 K에 대해 각자를 시작점으로 하는 그래프 탐색을 해줘야 모든 간선을 구할 수 있다.

 

 # 코드

import sys, collections

# find : 정점 x의 부모 노드를 리턴하는 함수
def find(x):
    if x == parent[x]:
        return x
    parent[x] = find(parent[x])
    return parent[x]

# merge : 정점 y의 부모를 정점 x의 부모로 변경하는 함수
def merge(x, y):
    x = find(x)
    y = find(y)
    parent[y] = x

# 입력부
n, m = map(int, sys.stdin.readline().split())
arr = [list(sys.stdin.readline().rstrip()) for _ in range(n)]

# parent : 정점 x의 부모 노드 번호를 저장하는 배열. 초기값으로는 자기 자신이다
parent = [i for i in range(m + 1)]
q = collections.deque()

# edge : 간선 리스트
edge = []

# cnt : 정점 번호
# to_node : 현재 위치 (x, y)를 key, 그 때 해당하는 정점 번호를 value로 갖는 dictionary
cnt = 0
to_node = dict()

# temp : 2차원 배열에서 S 혹은 K의 위치를 저장하는 배열
temp = []

for i in range(n):
    for j in range(n):
        if arr[i][j] == 'S' or arr[i][j] == 'K':
            # 현재 위치에 대한 정점 번호를 to_node에 저장 후 정점 번호 1증가
            to_node[(i,j)] = cnt
            cnt += 1
            temp.append((i, j))

dx = [0,1,0,-1]
dy = [1,0,-1,0]
for k in temp:
    q.append((k[0], k[1]))
    # check : 시작점에서 임의의 위치 (x, y)로 가기 위한 최소 거리를 저장하는 2차원 배열
              -1이라면 시작점에서 갈 수 없음을 의미
    check = [[-1] * n for _ in range(n)]
    check[k[0]][k[1]] = 0
    # BFS
    while q:
        x, y = q.popleft()
        for i in range(4):
            nx, ny = x + dx[i], y + dy[i]
            if 0 <= nx < n and  0 <= ny < n:
                if arr[nx][ny] != '1':
                    if check[nx][ny] == -1:
                        check[nx][ny] = check[x][y] + 1
                        q.append((nx, ny))
                        
    # 모든 정점에 대해 갈 수 있는 정점을 간선 리스트에 저장
    for i in temp:
        if check[i[0]][i[1]] > 0:
            edge.append((check[i[0]][i[1]], to_node[k], to_node[i]))

# 간선의 가중치 별 정렬 후 MST
edge.sort()
mst = 0
res = 0
for i in edge:
    # m - 1이 아닌 m은 이유는 총 정점이 m개의 K와 1개의 S이기 때문
    if res == m:
        break
    # 만일 부모가 다르다면 (사이클 형성이 되지 않는다면) union
    if find(i[1]) != find(i[2]):
        merge(i[1], i[2])
        res += 1
        mst += i[0]
        
# MST 형성 여부에 따른 정답 출력
print(mst if res == m else -1)