[Python/파이썬] 백준 1238번: 파티

문제

풀이

다익스트라 알고리즘을 사용하면 된다.

단, 학생들이 마을로 돌아가는 경우도 구해야 하는데, 이때의 간선은 주어진 입력의 시작점과 끝점을 서로 바꾸면 된다.

 

코드

from sys import stdin
import heapq

input = stdin.readline

n, m, x = map(int, input().split())

toX = [[] for _ in range(n + 1)] #집에서 파티 마을로 가는 최단 경로
toHome = [[] for _ in range(n + 1)] #파티 마을에서 집으로 가는 최단 경로

for _ in range(m):
      a, b, t = map(int, input().split())
      toX[a].append((b, t))
      toHome[b].append((a, t))

def dijk(graph, start):
      distance = [1e9] * (n + 1)
      q = []
      heapq.heappush(q, (0, start))
      distance[start] = 0
      while q:
            dist, now = heapq.heappop(q)
            if distance[now] < dist:
                  continue
            for i in graph[now]:
                  cost = dist + i[1]
                  if cost < distance[i[0]]:
                        distance[i[0]] = cost
                        heapq.heappush(q, (cost, i[0]))
      return distance[1:]

print(max([a+b for a, b in zip(dijk(toX, x), dijk(toHome, x))]))

 

마무리

맨 처음에는 시작점과 끝점을 바꿀 생각을 하지 못해 i (1<= i <=n) 노드의 distance[x] 값을 구한 뒤, 그 값을 x 노드의 distance[i]에 더해주었다. 문제는 풀 수 있으나 아직 실력이 많이 부족함을 느꼈다...