[백준] 1238번 파티(다익스트라 알고리즘)(Python - 파이썬)

백준 1238번 문제입니다. (solved.ac)기준 골드 3 문제입니다.

문제 접근 

n개의 숫자로 구분된 각각의 마을(노드)에 한 명의 학생이 살고 있습니다.m개의 단 방향 도로(간선)가 있고 x번 마을에 모여서 파티를 한다고 합니다. 학생들이 x번 마을에 갔다가 다시 돌아오는 최단 경로 중 가장 오래 걸리는 학생의 소요시간을 출력하면 됩니다.

최단 경로를 구해야하는 문제라서 다익스트라 알고리즘을 사용하여 접근하였습니다.

 

start라는 리스트를 만들어서 start[4]라면 4번째 마을에서 다른 모든 노드로 향하는 최단 경로를 저장하도록 하였습니다.

ex) 2번 마을에서 x번 마을로 갔다가 돌아오는 최단 경로는  start[2][x] + start[x][2] 입니다.

start[2][x]는 2번 마을에서 x번 마을로 가는 최단 경로, start[x][2]는 x번 마을에서 2번 마을로 가는 최단 경로

 

ans라는 리스트를 만들어 1번 마을부터 n번 마을까지 왕복 최단 경로들을 저장하도록 하였습니다.

ex) ans[1]은 1번 마을에서 x번 마을을 갔다오는 최단 경로 = start[1][x] + start[x][1]

그후 max(ans)를 출력하여 최단 경로 중 최댓값을 출력해주었습니다.

정답 코드

import heapq
INF = int(1e9)

# n은 노드 m은 간선 수 x는 도착 점
n, m, x = map(int, input().split())

graph = [[] for _ in range(n+1)]

start = [0] * (n+1)
ans = [0] * (n+1)

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

def dijkstra(start) :

    distance = [INF] * (n + 1)
    q = []
    distance[start] = 0

    heapq.heappush(q, (0, start))

    while q:

        dis, now = heapq.heappop(q)

        if distance[now] < dis :
            continue

        for i in graph[now] :
            cost = dis + i[1]

            if distance[i[0]] > cost :
                distance[i[0]] = cost
                heapq.heappush(q, (cost, i[0]))

    return distance

for i in range(1, n+1) :

     start[i] = dijkstra(i)

for i in range(1, n+1) :

    ans[i] = start[i][x] + start[x][i]

print(max(ans))