9370. 미확인 도착지 (Python)

 

9370번: 미확인 도착지

 

풀이

 다익스트라를 사용하는 문제인데 문제에서 주어진 g와 h 교차로 사이에 있는 도로를 지나가는 길을 포함하는 최소 경로가 있는지에 대해서 추가적으로 판단하는 작업이 필요하다. 보통 우선순위 큐에 넣을 때 (출발점에서의 거리, 다음 출발 노드)와 같은 형태로 넣어주는 데 이렇게 해서는 특정 길을 지났는지를 판단할 수 없다. 추가적으로 g와 h 도로를 지나간 경로인지를 판단하기 위한 플래그 값이 필요하다.

 

 위의 그림을 보면 1번에서 4번으로 가는 최소 경로 길이는 7이 된다. 문제는 1-2-4로 가는 경로와 1-3-4로 가는 경로가 모두 7의 길이를 가지기 때문에 일반적인 다익스트라로는 3-4 도로를 지난 경로를 찾았음이 보장되지 않는다. 이를 위해 우선순위 큐에 (출발점에서의 거리, g-h 도로를 지났는지 여부, 다음 출발 노드)와 같은 형태로 값을 넣는다. g-h를 지나지 않았다면 0을 넣고 g-h를 지났다면 -1을 넣어 출발점에서의 거리가 같다면 g-h를 지난 경로가 우선순위 큐에서 먼저 꺼내지도록 보장하는 것이다.

 

 또한 거리를 저장하는 distance 배열에 더불어 g-h를 지났는지를 기록하는 배열을 만들어 해당 지점까지의 최소 경로에서 g-h 도로를 지났는지를 기록하도록 한다.

 

코드

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


def dijkstra(start):
    q = []

    # (출발점에서의 거리, g-h 도로를 지났는지에 대한 플래그, 현재 노드)
    # 플래그가 0일 경우 -> g-h 도로를 지나지 않음, 플래그가 -1일 경우 -> g-h 도로를 지남
    heapq.heappush(q, (0, 0, start)) 

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

        # 출발점에서의 거리가 같을 경우 플래그가 -1인(g-h 도로를 지나는 경로)가 먼저 힙에서 꺼내지기 때문에
        # distance[now]와 dist가 같은 경우도 continue를 해준다.
        if distance[now] <= dist: 
            continue
        
        # 출발점에서 now까지의 거리와 g-h 도로를 지났는지에 대한 기록을 해준다.
        distance[now] = dist
        includeGH[now] = ghFlag

        for (next, nextDist) in roads[now]:
            cost = dist + nextDist
            
            # 보통 다익스트라는 <를 통해 비교하지만 해당 문제는 특정 도로를 지났는지를 판단해야 한다.
            # 그렇기 때문에 거리가 같은 경우도 힙에 넣어주도록 <= 비교 연산자를 사용한다.
            # 예를 들어, 처음에는 g-h 도로를 지나지 않은 거리가 7인 경로를 힙에 넣었는데
            # 추후에 g-h 도로를 지난 거리가 7인 경로를 찾을 수 있다.
            if cost <= distance[next]:
                # g-h 도로를 지났을 경우 플래그를 -1로 설정하여 힙에 넣어준다.
                if (now, next) == (g, h) or (now, next) == (h, g):
                    heapq.heappush(q, (cost, -1, next))
                else: # 그렇지 않을 경우 기존의 플래그 값을 넣어준다.
                    heapq.heappush(q, (cost, ghFlag, next))



if __name__ == '__main__':
    T = int(input())

    for _ in range(T):
        n, m, t = map(int, input().split())
        s, g, h = map(int, input().split())
        roads = [[] for _ in range(n+1)]
        distance = [INF] * (n+1)
        includeGH = [0] * (n+1) # 최소 경로가 g-h 도로를 지났는지 여부. -1이면 지났음을 뜻한다.
        destinations = set() 

        for _ in range(m):
            a, b, d = map(int, input().split())
            roads[a].append((b, d))
            roads[b].append((a, d))

        for _ in range(t):
            destinations.add(int(input()))

        dijkstra(s)

        # g-h 도로를 지나는 최소 경로를 가진 도시들을 trueIncludeGH 셋에 넣어준다.
        # 그 후 목적지 후보 set과 교집합 연산을 통해서 문제의 답을 구한다.
        trueIncludeGH = set([x for x in range(n+1) if includeGH[x] == -1])
        resList = sorted(list(destinations.intersection(trueIncludeGH)))
        
        for res in resList:
            print(res, end=' ')
        print()