백준 1238: 파티 - 다익스트라 (Python/파이썬)

다익스트라

최단 거리 알고리즘. 다만 간선의 가중치가 음수일 땐 사용할 수 없음.
0. 출발 노드 설정
1. 출발 노드부터 각 노드까지의 최소 비용 저장(미방문 - 무한대로 저장)
2. 방문하지 않은 노드 중에 제일 비용 적은 노드 선택
3. 해당 노드를 거쳐 다른 노드로 가는 경우 고려해서 최소 비용 갱신

2 - 3 반복

1. 첫 제출

import sys
import heapq

input = sys.stdin.readline

def dijkstra(start, target=None):
    distance = [float('inf')] * (n+1)
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue

        for node in graph[now]:
            v, e = node
            cost = dist + v
            if distance[e] > cost:
                distance[e] = cost
                heapq.heappush(q, (cost, e))

    if target:
        return distance[target]
    else:
        return distance

n, m, x = map(int, input().split())
graph = [[] for _ in range(n+1)]

for _ in range(m):
    start, end, time = map(int, input().split())
    graph[start].append((time, end))

cost = dijkstra(x)
for i in range(1, n+1):
    cost[i] += dijkstra(i, x)

print(max(cost[1:]))

단방향 그래프라 가는 길, 오는 길을 각각 다익스트라를 돌려야 하는데 처음엔 각 노드에서 x 가는 길을 for문으로 다익스트라 함수 n번 돌림.

2. 개선 코드

애초에 input 받을 때 역방향 그래프에 저장해 각 노드에서 해당 노드까지 오는 거리를 계산하도록 변경.
시간 엄청 줄었는데 메모리는 거기서 거기.

import sys
import heapq

input = sys.stdin.readline

def dijkstra(start, graph):
    distance = [float('inf')] * (n+1)
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0

    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue

        for node in graph[now]:
            cost = dist + node[0]
            if distance[node[1]] > cost:
                distance[node[1]] = cost
                heapq.heappush(q, (cost, node[1]))

    return distance

n, m, x = map(int, input().split())
g = [[] for _ in range(n+1)]
reverse_g = [[] for _ in range(n+1)]

for _ in range(m):
    start, end, time = map(int, input().split())
    g[start].append((time, end))
    reverse_g[end].append((time, start))

togo = dijkstra(x, g)
tocome = dijkstra(x, reverse_g)

ans = 0
for i in range(1, n+1):
    if i == x:
        continue
    if ans < togo[i] + tocome[i]:
        ans = togo[i] + tocome[i]

print(ans)