[백준 파이썬] 1238 파티

1238번: 파티

문제

 

 

---

아이디어

 그래프의 최단 경로를 구하여 코스트를 계산하는 문제이므로 다익스트라 알고리즘을 사용하였다.

---

풀이

heapq를 이용해 다익스트라를 구현하였다.

def dijkstra(start):
    q = []
    distance = [sys.maxsize] * (N + 1)
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)

        if distance[now] < dist:
            continue
        for nodeIndex, nodeCost in data[now]:
            cost = dist + nodeCost

            if distance[nodeIndex] > cost:
                distance[nodeIndex] = cost
                heapq.heappush(q, (cost, nodeIndex))
    return distance

 

 

 

result = 0
for i in range(1, N + 1):
    go = dijkstra(i)
    back = dijkstra(X)
    result = max(result, go[X] + back[i])

 

파티에 갈때의 최소경로와 돌아올 때의 최소경로를 구해야하므로 다익스트라를 2번사용하였고, 각각의 경우를 더한값의 가장 큰 코스트를 얻어 출력하여 해결하였다.

 

---

코드

import heapq
import sys
N, M, X = map(int, input().split())

data = [[] for _ in range(N+1)]

for i in range(M):
    a, b, c = map(int, input().split())
    data[a].append((b, c))


def dijkstra(start):
    q = []
    distance = [sys.maxsize] * (N + 1)
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)

        if distance[now] < dist:
            continue
        for nodeIndex, nodeCost in data[now]:
            cost = dist + nodeCost

            if distance[nodeIndex] > cost:
                distance[nodeIndex] = cost
                heapq.heappush(q, (cost, nodeIndex))
    return distance


result = 0
for i in range(1, N + 1):
    go = dijkstra(i)
    back = dijkstra(X)
    result = max(result, go[X] + back[i])


print(result)