[백준 / Python] #23973 두 단계 최단 경로 1

문제

알고리즘

다익스트라를 활용하여 Y노드를 포함해서 Z까지의 최단경로와, Y노드를 지나지 않고 Z까지의 최단경로를 구하는 문제이다.

Y노드를 무조건 지나는 경로는 X to Y + Y to Z의 최단경로를 서로 더해주면 된다.

Y노드를 지나지 않는 최단경로를 구할때는, 다익스트라를 반복하면서 Y노드를 방문할때 해당 계산을 무시하고 건너뛰도록 설정해주면 된다.

코드

import sys
import heapq
input = sys.stdin.readline

INF = int(1e9)

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

for _ in range(m):
    u, v, w = map(int, input().split())
    graph[u].append((v, w))

x, y, z = map(int, input().split())

def dijkstra(start):
    distance = [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 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

# x to y to z
dist_from_x = dijkstra(x)
dist_from_y = dijkstra(y)

path_x_to_y = dist_from_x[y]
path_y_to_z = dist_from_y[z]

path_with_y = -1
if path_x_to_y != INF and path_y_to_z != INF:
    path_with_y = path_x_to_y + path_y_to_z

# x to z, not visiting y
distance_no_y = [INF] * (n + 1)
q_no_y = []
heapq.heappush(q_no_y, (0, x))
distance_no_y[x] = 0

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

    if distance_no_y[now] < dist:
        continue

    for i in graph[now]:
        neighbor_node = i[0]
        # if visiting node is y, ignoring
        if neighbor_node == y:
            continue

        cost = dist + i[1]

        if cost < distance_no_y[neighbor_node]:
            distance_no_y[neighbor_node] = cost
            heapq.heappush(q_no_y, (cost, neighbor_node))

path_without_y = distance_no_y[z]
if path_without_y == INF:
    path_without_y = -1

print(path_with_y, path_without_y)