[백준 1944번] 복제 로봇 (Python3)

import sys,collections,heapq
input = sys.stdin.readline
dy = [1,-1,0,0]
dx = [0,0,1,-1]

def find():
  for y in range(N):
    for x in range(N):
      if board[y][x]=="S" or board[y][x]=="K":
        Index[(y,x)] = len(location)
        location.append((y,x))

def BFS(y,x,now):
  cnt = 0
  visited = [[0]*N for i in range(N)]
  dq = collections.deque([(y,x,0)])
  while dq:
    y,x,d = dq.popleft()
    if visited[y][x]:
      continue
    visited[y][x] = 1
    if board[y][x] != "0":
      next = Index[(y,x)]
      graph[now][next] = d
      cnt += 1
    for i in range(4):
      y1,x1 = y+dy[i],x+dx[i]
      if not visited[y1][x1] and board[y1][x1] != "1":
        dq.append((y1,x1,d+1))
  if cnt == M+1:
    return True

def MST():
  for now in range(M+1):
    if not BFS(*location[now],now):
      return -1
  result = 0
  visited = [0]*(M+1)
  hq = []
  for i in range(M+1):
    heapq.heappush(hq,(graph[0][i],i))  
  for _ in range(M+1):
    while hq:
      w,next =  heapq.heappop(hq)
      if not visited[next]:
        visited[next] = 1
        result += w
        for i in range(M+1):
          heapq.heappush(hq,(graph[next][i],i))
        break
  return result

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

board = [input().strip() for i in range(N)]

location,Index = [],{}
find()
graph = [[0]*(M+1) for i in range(M+1)]
print(MST())

MST (최소신장트리) 알고리즘으로 해결하였다.

처음에 문제를 읽고 열쇠를 찾는 시간을 구하는 문제라고 착각해서 "뭐야 그냥 BFS하면 되는거 아니야?" 했는데, 다시 읽어보니까 로봇의 이동 거리를 구하는 문제였다.

 

따라서,

1. BFS로 시작지점+열쇠의 서로간의 거리를 기록한다.

2. MST 알고리즘으로 시작지점+열쇠를 모두 잇는 최소거리를 구한다.

로 간단하게 해결할 수 있다.

 

 

오늘의 교훈) 문제를 정확하게 파악하자