[백준] 1238 파티 python3

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

1238번: 파티

 

코드

import heapq
from sys import stdin
input = stdin.readline
INF = float("inf")


def dijkstra(start, end):
    dist = [INF for _ in range(N+1)]
    dist[start] = 0

    q = []
    heapq.heappush(q, (0, start))

    while q:
        time, village = heapq.heappop(q)

        for ntime, nviliage in graph[village]:
            ntime += time
            if dist[nviliage] > ntime:
                dist[nviliage] = ntime
                heapq.heappush(q, (ntime, nviliage))

    return dist

N, M, X = map(int, input().split(" "))
graph = [[] for _ in range(N+1)]
for i in range(M):
    v, u, t = map(int, input().split(" "))
    graph[v].append((t, u))

# 파티 장소로 가기
max_time = -1
transfer_times = [[0] for _ in range(N+1)]
for i in range(1, N+1):
    if i == X:
        continue
    transfer_times[i] = dijkstra(i, X)[X]

comeback_times = dijkstra(X, None)
# 집으로 돌아가기
for i in range(1, N+1):
    if i == X:
        continue
    transfer_times[i] += comeback_times[i]
    max_time = max(transfer_times[i], max_time)

print(max_time)

 

해설

다익스트라를 두 번 돌려서 푸는 문제이다.

기존에 다익스트라를 풀던대로 그래프를 dict로 하고 다음 탐색대상 노드를 따로 찾으며 했더니 시간초과가 계속 걸렸다.

힙을 활용해서 우선순위큐로 풀면 시간이 훨씬 단축되는 듯 하다.