[백준 1944번] 복제로봇(파이썬)

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

1944번: 복제 로봇

이 문제를 처음 풀때는

이렇게 생각했습니다.

 

1. 로봇을 기준으로 key까지의 거리를 구한다

 

2. 각 key들간의 거리를 구한다.

 

1까지는 그렇다치더라도

2에서 시간초과가 엄청날 것으로 보였습니다.

key의 갯수가 250개가 된다하면

1->2~250번째 key까지의 거리

2->2를 제외한 나머지 키들까지의 거리

3->.....

이렇게 구하다보면....

 

역시나 시간초과판정을 받고

고민을 하다가

구글링해보니

이 문제를 풀기 위해서는

크루스칼 알고리즘을 알아야했습니다.

https://techblog-history-younghunjo1.tistory.com/262

[Python] 최소 신장 트리를 찾는 크루스칼(Kruskal) 알고리즘

이 글이 크루스칼 알고리즘에 대한 설명입니다.

 

여기서 heap의 push할 때 어떻게 넣는지 알면

이 문제를 좀 더 잘 이해할 수 있는데요

heappush를 할 때는

앞의 값을 기준으로

자연스럽게 정렬이 됩니다.

예를들어보면

8 a b /3 a c/9 c d 이렇게 push했다가

나중에 pop을 할 때는

3 a c / 8a b/9 c d 순으로 출력이 됩니다.

처음에 이 문제에 대한 알고리즘을 볼 때

이부분이 이해가 안되서 당황했는데

heap의 동작을 다시 생각해보니

감이 잡혔습니다.

 

#백준 1944번 복제로봇
from collections import deque
from heapq import heappop,heappush
from itertools import combinations

N,M=map(int,input().split())

room=[input().rstrip() for _ in range(N)]

dx=[0,0,1,-1]
dy=[1,-1,0,0]

node=[]
node_cnt=0
for i in range(N):
    for j in range(N):
        if room[i][j]=='S' or room[i][j]=='K':
            node.append([i,j])

answer=0


#1. 시작점을 기준으로 bfs해준 걸 담아본다
#시작점에서 로봇이 복제가 되는 케이스
def bfs(a,b):
    sx,sy=node[a]
    ex,ey=node[b]
    q = deque()
    q.append([0,sx,sy])
    visited = [[False] * N for _ in range(N)]
    visited[sx][sy] = True
    while q:
        cnt, x, y = q.popleft()
        if x == ex and y == ey:
            return cnt

        for d in range(4):
            nx = x + dx[d]
            ny = y + dy[d]

            if nx < 0 or nx >= N or ny < 0 or ny >= N:
                continue

            if room[nx][ny] == '1':
                continue

            if visited[nx][ny] == True:
                continue
            visited[nx][ny] = True
            q.append([cnt+1,nx,ny])
    return 0

def find(u):
    if u!=p[u]:
        p[u]=find(p[u])
    return p[u]

def union(a,b):
    r1=find(a)
    r2=find(b)
    p[r2]=r1

p=[i for i in range(M+1)]
way=[]

for a,b in combinations(range(M+1),2):
    w=bfs(a,b)
    if w==0:
        print(-1)
        exit(0)

    heappush(way,[w,a,b])
edge=0


while True:

    if edge==len(node)-1:
        break
    if not way:

        print(-1)
        exit(0)
    w,a,b=heappop(way)
    if find(a)!=find(b):
        union(a,b)
        answer+=w
        edge+=1


print(answer)

 크루스칼 알고리즘에 대해서 학습하기 

아주 좋은 문제입니다.

문제 자체를 곰곰히 곱씹어보기 좋습니다.