[백준] 1238번: 파티

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

1238번: 파티

풀이

출발지와 목적지 사이의 최단 왕복 시간을 구하는 전형적인 최단 경로 문제이다. 다만 한 정점에서만 출발하는 것이 아니라, N개의 정점 각각에서 출발했을 때의 최단 경로를 구하고, 그중에서 최대 값을 뽑아야 한다.

 

처음엔 문제를 대충 읽고, 단지 모든 정점에서의 최단 경로를 구해야한다는점 만을 고려해 플로이드 워셜 알고리즘으로 접근했다.

import sys

N, M, X = map(int, sys.stdin.readline().rstrip().split()) 

graph = [[inf]*(N+1) for _ in range(N+1)]

for i in range(1, N+1):
  for j in range(1, N+1):
    if i==j:
      graph[i][j] = 0

for _ in range(M):
  a, b, c = map(int, sys.stdin.readline().rstrip().split())
  graph[a][b] = c

for k in range(1, N+1):
  for i in range(1, N+1):
    for j in range(1, N+1):
      graph[i][j] = min(graph[i][j], graph[i][k] + graph[k][j])

max_distance = 0
for i in range(1, N+1):
  max_distance = max(graph[i][X] + graph[X][i], max_distance)

print(max_distance)

문제 예제의 답은 얻을 수 있었지만 시간 초과로 오답 처리 되었다. 간선 M이 최대 10,000개까지 들어오기 때문인데, 플로이드 워셜 알고리즘의 시간 복잡도는 O(N³)이므로 굉장히 느려, 보통 N이 100 이상인 경우엔 사용할 수 없다.

 

시간 초과와 더불어, 문제를 다시 읽어보니 목적지는 정점 X 한 곳이기 때문에, 다익스트라 알고리즘을 사용하는 것이 더 적합하다는 것을 깨닳았다. 다익스트라 알고리즘을 사용한 풀이는 아래와 같다.

import sys
import heapq
from math import inf
from collections import defaultdict

N, M, X = map(int, f.readline().rstrip().split()) 

graph = defaultdict(list)

for _ in range(M):
  a, b, c = map(int, f.readline().rstrip().split())
  graph[a].append([c, b])

def djikstra(start, end):
  global graph

  Q = [[0, start]]
  dist = [inf]*(N+1)
  dist[start] = 0

  while Q:
    time, node = heapq.heappop(Q)

    if time > dist[node]:
      continue
    for c, b in graph[node]:
      if dist[b] > c + time:
        dist[b] = c + time
        heapq.heappush(Q, [c+time, b])

  return dist[end]

print(max([djikstra(i, X) + djikstra(X, i)  for i in range(1, N+1)]))