백준 / 미확인 도착지 / python

import sys
input = sys.stdin.readline

from heapq import heappop, heappush

T = int(input())
def dijktra(n, s, graph):
    dist = [1e9] * (n+1)
    dist[s] = 0

    h = []

    heappush(h, (0, s))

    while h:
        dis, node = heappop(h)
        if dist[node] < dis:
            continue

        for next_node, next_dis in graph[node]:
            d = dis + next_dis
            if dist[next_node] > d:
                dist[next_node] = d
                heappush(h, (d, next_node))

    return dist

def check(s, g, h, o_dij, g_dij, h_dij, r):

    if o_dij[g] + g_dij[h] + h_dij[r] == o_dij[r]:
        return True
    
    if o_dij[h] + h_dij[g] + g_dij[r] == o_dij[r]:
        return True

    return False

for _ in range(T):
    # 교차로, 도로, 목적지 후보의 개수
    n, m, t = [int(v) for v in input().split()]

    # 출발지, 지나간 교차로 (g,h)
    s, g, h = [int(v) for v in input().split()]

    graph = [[] for _ in range(n+1)]

    for i in range(m):
        a, b, d = [int(v) for v in input().split()]
        graph[a].append((b,d))
        graph[b].append((a,d))
    
    r = []

    for i in range(t):
        r.append(int(input()))

    o_dij = dijktra(n, s, graph)
    g_dij = dijktra(n, g, graph)
    h_dij = dijktra(n, h, graph)

    result = []
    for i in sorted(r):
        if check(s,g,h, o_dij, g_dij, h_dij, i):
            result.append(i)
    print(*result)

다익스트라 문제 중 특정한 최단 경로와 굉장히 유사한 문제

예술가는 반드시 시작점으로부터 최단경로로 이동해야 하며 g, h 사이에 있는 도로를 반드시 지나야한다.

즉 g, h 노드를 지나면서 목적지 후보에 도착한 값이 시작점에서 목적지 후보에 도착한 값과 똑같으면 가능한 경우이다. 이 경우를 판단하는 것을 check 함수로 작성