[Python/파이썬] 백준 1944번 복제 로봇

1944번: 복제 로봇

 

코드

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

n, m = map(int, input().split())
maze = []
points = []
for i in range(n):
    row = list(input())
    maze.append(row)
    for j in range(n):
        if row[j] == 'S' or row[j] == 'K':
            maze[i][j] = len(points)
            points.append((i, j))

# BFS
graph = []


def bfs(idx):
    dx = [0, 0, -1, 1]
    dy = [-1, 1, 0, 0]

    x, y = points[idx]
    q = deque([(x, y, 0)])
    visited = [[False]*n for _ in range(n)]
    visited[x][y] = True
    while q:
        x, y, cnt = q.popleft()
        for i in range(4):
            nx = x+dx[i]
            ny = y+dy[i]
            if 0 <= nx < n and 0 <= ny < n and not visited[nx][ny] and maze[nx][ny] != '1':
                q.append((nx, ny, cnt+1))
                visited[nx][ny] = True
                if maze[nx][ny] != '0':
                    graph.append((cnt+1, idx, maze[nx][ny]))


for i in range(len(points)):
    bfs(i)

# MST(Kruskal Algorithm)
parent = [i for i in range(len(points))]


def find_parent(x):
    if x != parent[x]:
        parent[x] = find_parent(parent[x])
    return parent[x]


def union_parent(a, b):
    a = find_parent(a)
    b = find_parent(b)

    if a == b:
        return False

    if a < b:
        parent[b] = a
    else:
        parent[a] = b
    return True


graph.sort()
ans = 0
num = 0
for dist, s, e in graph:
    if union_parent(s, e):
        ans += dist
        num += 1
    if num >= len(points)-1:
        break

print(ans if num == len(points)-1 else -1)

 

이 문제는 로봇의 시작점 S와 키가 놓인 위치 K를 각각 노드로 보고 해당 노드들을 모두 이어서 MST로 만들었을 때의 가중치 총합을 구하면 된다.

 

우선 MST를 만드려면 각 노드 간의 거리(가중치)를 알아야 한다. 'S'와 'K'의 위치를 points 리스트에 저장해둔 뒤, BFS를 이용하여 points에 있는 각 칸들 간의 거리를 구하여 graph 배열에 [거리, x좌표, y좌표] 형식으로 저장한다.

 

그리고 길이가 짧은 간선부터 하나씩 선택하여 MST에 포함시키는 크루스칼 알고리즘을 사용하기 위해 graph를 거리 오름차순으로 정렬한다.

그런 뒤에 길이가 짧은 간선부터 하나씩 선택한다. 이때 MST는 사이클을 만들면 안되기 때문에 사이클을 형성하는 간선은 선택하지 않는다. MST의 간선의 개수는 노드 개수-1개이므로 선택한 간선이 len(points)-1개가 될 때까지 간선을 선택한다.

만약 graph의 모든 간선들을 살펴보며 선택해도 len(points)-1개만큼 선택할 수 없다면 MST를 만들 수 없는 경우이므로 이 경우에는 -1을 출력한다. 그렇지 않은 경우는 선택한 간선들의 가중치의 총합을 출력해준다.