백준 9370번 : 미확인 도착지 in Python

9370번: 미확인 도착지

 

www.acmicpc.net

이 문제는 후보로 주어진 여러 목적지 중 최단 거리로 이동 시 특정 간선을 지나는 목적지를 찾는 문제로, 이전의 1504번 '특정한 최단 경로'와 유사한 문제이나 이 문제는 특정 정점이 아닌 특정 간선을 지나야 한다는 점에서 약간 차이가 있다.

1504번: 특정한 최단 경로

 

www.acmicpc.net

처음에는 이 문제를 풀 때 우선 각 목적지에 대해 출발점에서 최단 거리로 이동하는 경로를 찾아 그 경로에 특정 간선이 포함되는지 검사하려고 했다. 그래서 기존의 Dijkstra 알고리즘을 구현하는 함수에 root 배열을 추가해 최단 거리로 이동하는 경로를 저장했다. 그러나 이렇게 구현한 코드를 제출한 결과 틀렸다는 결과가 나왔다.

그래서 왜 그런지 생각해본 결과 특수한 경우로 목적지까지 최단 경로가 두 개 이상 존재할 때 위와 같은 방식으로 구현하면 그 경로들 중 하나만 알 수 있다는 사실을 알았다. 그래서 위와 같은 방법이 아닌 다른 방법으로 특정 간선을 지나는 지를 검증하기로 했다.

우선 Dijkstra 알고리즘을 이용해 출발점에서 각 후보 목적지까지 최단 경로의 길이를 구한다. 그 다음에 특성 간선이 점 g와 점 h를 잇는 선이라고 할 때, (출발점에서 g까지의 최단 경로의 길이 + 특정 간선의 길이 + g에서 목적지까지의 최단 경로의 길이)와 (출발점에서 h까지의 최단 경로의 길이 + 특정 간선의 길이 + h에서 목적지까지의 최단 경로의 길이)를 비교해 더 작은 값을 구한다. 그리고 앞에서 구한 출발점에서 목적지까지의 최단 경로의 길이와 비교해 만약 같은 경우 특정 간선을 포함한 경로가 출발점에서 목적지까지의 최단 경로 중 하나라는 결론을 얻을 수 있다. 즉, 해당 목적지는 특정 간선을 지나는 경우 최단 거리로 이동할 수 있다는 사실을 알 수 있다. 이러한 검증을 후보 목적지에 대해 모두 실행해 검증에 성공한 목적지를 모아 오름차순으로 정렬한 뒤 출력하면 된다.

이를 코드로 구현하면 다음과 같다.

import sys
import heapq
def Dijkstra(v):
    global n
    dijk = [sys.maxsize for _ in range(n+1)]
    visit = [False for _ in range(n+1)]
    dijk[v] = 0
    heap = []
    heapq.heappush(heap, [0, v])
    while len(heap) > 0:
        [d, u] = heapq.heappop(heap)
        if visit[u] == True: continue
        visit[u] = True
        for [v, w] in graph[u]:
            if dijk[v] > dijk[u] + w:
                dijk[v] = dijk[u] + w
                heapq.heappush(heap, [dijk[v], v])
    return dijk
T = int(input())
for _ in range(T):
    n, m, t = map(int, input().split())
    s, g, h = map(int, input().split())
    graph = [[] for _ in range(n+1)]
    for _ in range(m):
        a, b, d = map(int, input().split())
        graph[a].append([b, d])
        graph[b].append([a, d])
    # 후보 목적지를 저장하는 배열
    des = []
    for _ in range(t):
        des.append(int(input()))
    dijk1 = Dijkstra(s)
    dijk2 = Dijkstra(g)
    dijk3 = Dijkstra(h)
    # g와 h를 잇는 간선의 길이를 구한다.
    e = sys.maxsize
    for [b, d] in graph[g]:
        if b == h: e = d
    res = []
    # 각 후보 목적지에 대해 최단 경로가 간선을 포함하는지 확인한다.
    for i in range(t):
        if dijk1[des[i]] == min(dijk1[g] + e + dijk3[des[i]], dijk1[h] + e + dijk2[des[i]]):
            res.append(des[i])
    # 오름차순으로 정렬한다.
    res.sort()
    res = map(str, res)
    print(" ".join(res))