Coding 공부

import sys
import heapq
from collections import deque
def input():
    return sys.stdin.readline()

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

arr = [list(input()) for _ in range(N)]

start = [0,0,0]
node_list = []
total_cnt = 0
for x in range(N):
    for y in range(N):
        if arr[x][y] == 'S':
            start = [0,x,y]
            arr[x][y] = '0'



heapq.heappush(node_list,start)
INF = float('inf')
visited = [[True for _ in range(N)] for _ in range(N)]
distance = [[INF for _ in range(N)] for _ in range(N)]
result = 0
dx = [-1,1,0,0]
dy = [0,0,-1,1]
cnt = 0
while node_list:
    cur_dis,x,y, = heapq.heappop(node_list)
    if not visited[x][y]:
        continue
    if cur_dis > distance[x][y]:
        continue
    result += cur_dis
    cnt += 1
    visited[x][y] = False

    temp_visited = [[True for _ in range(N)] for _ in range(N)]
    queue = deque()
    queue.append((x,y,0))
    temp_visited[x][y] = False
    while queue:
        qx,qy,dis = queue.popleft()

        for i in range(4):
            nx = qx + dx[i]
            ny = qy + dy[i]

            if temp_visited[nx][ny] and arr[nx][ny] != '1':
                temp_visited[nx][ny] = False
                queue.append((nx,ny,dis+1))
                if arr[nx][ny] == 'K' and distance[nx][ny] > dis+1:
                    distance[nx][ny] = dis + 1
                    heapq.heappush(node_list,(dis+1,nx,ny))
if cnt<M+1:
    print(-1)
else:
    print(result)

 이 문제는 MST와 너비 우선 탐색을 활용한 문제였다.

일단 경계선은 전부 1로 되어있으므로, 경계체크를 해줄 필요는 없다.

문제를 푸는 아이디어는 다음과 같다. 

프림알고리즘을 이용한다.

그런데 프림 알고리즘을 활용할려면, 각 노드 사이의 간선의 정보를 알아야한다.

이 것은 BFS를 통해 로봇과 K들 사이의 간선들을 구하는 것이다.

즉, heap에 들어온 Node를 기준으로 다른 열쇠들 간의 간선의 정보를 찾으면 된다.

그래서 프림 알고리즘을 해서 나오는 노드를 기준으로 BFS를 돌리고, 열쇠이고, 현재 누적된 거리가 더 짧으면

거리를 갱신하고, 그 때의 위치 정보 및 거리를 heap에 넣어준다.

이런식으로 해서 프림알고리즘을 돌리면서 매번 BFS를 돌려서 해당 노드에서 다음노드로 갈수 있는 간선정보를 찾아

heap에 넣어주는 식으로 해주면 된다.

그리고 여기서 주의해야할 점은 모든 열쇠를 찾지 못하는 경우도 있다.

우리는 총 M + 1개의 노드를 가진 그래프에서 프림 알고리즘을 돌린것과 마찬가지이므로,

프림 알고리즘에서 방문한 노드의 수가 M+1개 보다 적으면 -1을 출력 해주며 된다.

import sys
import heapq

input = sys.stdin.readline

N = int(input())
M = int(input())
graph = [{} for i in range(N)]

for _ in range(M):
    a,b,c = map(int,input().split())
    if a != b:
        graph[a-1][b-1] = c
        graph[b-1][a-1] = c

INF = float('inf')
distance = [INF]*N
visited = [False] *N
distance[0] = 0
node_list = []
heapq.heappush(node_list,(0,0))
result = 0
while node_list:
    dis, node = heapq.heappop(node_list)
    if visited[node]:
        continue
    result += dis
    visited[node] = True
    for next_node in graph[node]:
        if distance[next_node] > graph[node][next_node]:
            distance[next_node] = graph[node][next_node]
            heapq.heappush(node_list,(distance[next_node],next_node))

print(result)

최소 스패닝 트리 문제이다. 프림알고리즘과 크루스칼 알고리즘이 있는데, 크루스칼 알고리즘에 아직 익숙하지 않아, 프림 알고리즘으로 구현했다.

heapq를 이용해서 구현했기 때문에, distance가 갱신될때에만 heapq.heappush를 해주는 것만 주의해주면 풀 수 있는 문제이다.