백준 1504번 문제(특정한 최단 경로) 파이썬(Python) 풀이 [로밍맨]

문제 링크

www.acmicpc.net/problem/1504

1504번: 특정한 최단 경로

 

정답 코드는 아래와 같습니다.

import sys
import math
import heapq
 
def dijkstra(graph, n, source, sink1, sink2):
  dist = [math.inf] * n
  pq = []
  dist[source] = 0
  heapq.heappush(pq, (0, source))
  while len(pq) != 0:
    cost, node = heapq.heappop(pq)
    for i in range(1, n):
      nextCost = dist[node] + graph[node][i]
      if nextCost < dist[i]:
        dist[i] = nextCost
        heapq.heappush(pq, (nextCost, i))
  return dist[sink1], dist[sink2]
 
def solve():
  n, e = map(int, sys.stdin.readline().rstrip().split())
  graph = [[math.inf] * (n+1) for _ in range(n+1)]
  for _ in range(e):
    a, b, c = map(int, sys.stdin.readline().rstrip().split())
    graph[a][b] = c
    graph[b][a] = c
  v1, v2 = map(int, sys.stdin.readline().rstrip().split())
  a, d = dijkstra(graph, n+1, 1, v1, v2)
  b, f = dijkstra(graph, n+1, v1, v2, n)
  c, e = dijkstra(graph, n+1, v2, n, v1)
  rst = min(a+b+c, d+e+f)
  if rst != math.inf:
    print(rst)
  else:
    print(-1)
 
solve()
 

 

728x90

 

전형적인 그래프 최단거리 문제로 가중치가 모두 양수인 것을 고려하여 다익스트라 알고리즘을 이용하여 풀 수 있는 문제입니다.

다익스트라는 기업 코딩테스트에서 매우 자주 출제되는 문제이기 때문에 거의 외우는 수준으로 숙달하시기 바라며,

priority queue 자료구조를 쉽게 사용할 수 있는 python이 아닌 다른 언어를 사용하시는 경우에는 priority queue 구현까지도 같이 외우셔야 실제 상황에서 문제를 풀 수 있습니다.

 

풀이는 아래 영상을 참고 바랍니다.

www.youtube.com/watch?v=oDvzkPNhG18  

 

저작권 라이선스: CC BY (출처만 표시하면 자유롭게 이용 가능)