[백준 17835번] 면접보는 승범이네 (Python3)

import sys
input = sys.stdin.readline
from heapq import *

N,M,K = map(int,input().split())
graph = [[] for i in range(N+1)]
for _ in range(M):
  a,b,w = map(int,input().split())
  graph[b].append((a,w))
data = [*map(int,input().split())]
check = [0]*(N+1); hq = []
for a in data:
  check[a] = 1
  for b,w in graph[a]:
    heappush(hq,(w,-b))
now,w = -1,0
for _ in range(N-K):
  while 1:
    w,now = heappop(hq)
    if check[-now]:
      continue
    break
  check[-now] = 1
  for next,w1 in graph[-now]:
    heappush(hq,(w+w1,-next))
print(-now);print(w)

간단한 프림 알고리즘 응용문제

K개의 면접장소는 이미 연결되어있다고 가정하고, 나머지 N-K개를 연결하면 된다. 이때 연결가중치는 w가 아니라 w+이전w 로 넣어주어야한다. (다익스트라처럼)

마지막에 연결된 노드를 출력하면 끝

 

이때 주의해야할 점은

1. 가중치가 같을 경우 가장 작은 번호를 출력해야한다. 나는 노드의 번호를 음수로 입력해서 가장 마지막에 가장 작은 번호가 정렬되도록 하는 방법을 이용하였다.

2. N=K일 수 있다. 이 경우 1,0을 출력해야한다.

 

 

오늘의 교훈) 프림 알고리즘은 유용하다.