BaekJoon 9370번 : 미확인 도착지 (python)

BaekJoon 9370번 : 미확인 도착지 (G2 25.185%)

이번 미확인 도착지 문제도 특정한 최단 경로와 매우 비슷한 문제이다.
특정 구간을 꼭 지나가야한다는 조건이 같고 여러 도착지가 있을 수 있다는 조건이 달린 문제이다.
문제에서는 주어진 조건인 상태에서 최단 경로인 상황인 도착지들을 출력하는 문제이다.

제일 처음에는 문제를 제대로 읽지 않아서 주어진 도착 리스트들의 최단경로들을 구하는 문제인 줄 알고 여러번 틀렸었다. 그 다음에는 다익스트라 알고리즘을 반복문에 넣어서 매번 구하는 비효율 때문에 시간복잡도가 너무 커져서 시간초과가 뜨기도 했었다.

• 결국 이 문제에서는 특정한 최단 경로를 구한 후 이 경로가 최단 경로일 경우 그 도착지를 출력하면 되는 문제이다.
• 특정한 최단 경로를 구하기 위해 입력받은 g, h와 시작지점인 s의 다익스트라 알고리즘을 구해서 리스트를 반환 받는다.
• 이렇게 되면 S,G,H에서 각각 시작한 최단 경로값들을 받을 수 있고 [S,G,H,E]or[S,G,E,H]가 [S,E]와 똑같은지 아닌지를 비교해 볼 수 있다.
• 만약 비교했을 때 같은게 존재한다면 정답이고 만약 두 경우의 수가 최단 경로의 경우의 수가 아니라면 답이 되지 못한다.

import sys,heapq
input = sys.stdin.readline
INF = 100000000

def dijkstra(start) :
    q = list()
    heapq.heappush(q,(start, 0))
    distance = [INF]*(N+1)
    distance[start] = 0
    while q :
        now, weight = heapq.heappop(q)

        if distance[now] < weight :
            continue
        for i in graph[now] :
            cost = i[1] + weight
            if distance[i[0]] > cost :
                distance[i[0]] = cost
                heapq.heappush(q,(i[0], cost))
    return distance

T = int(input())
for _ in range(T) :
    N, M, K = map(int,input().split())  #K는 목적지 후보의 개수
    S, G, H = map(int,input().split())  #s시작, g와 h를 무조건 지난다.
    graph = [[] for _ in range(N+1)]
    endcase = [0]*K
    for _ in range(M) :
        a, b, c = map(int,input().split())
        graph[a].append((b,c))
        graph[b].append((a,c))
    for i in range(K) :
        endcase[i] = int(input())
    result = list()
    s = dijkstra(S)
    g = dijkstra(G)
    h = dijkstra(H)
    
    for i in range(K) :
        a = s[G] + g[H] + h[endcase[i]]
        b = s[H] + h[G] + g[endcase[i]]
        if s[endcase[i]] == a or s[endcase[i]] == b :
            result.append(endcase[i])
    result.sort()
    print(*result)