백준 9370 - 미확인 도착지 (파이썬)

Answer

from sys import stdin
from heapq import heappush, heappop

input = stdin.readline

INF = int(1e9)

def dijkstra(start):
    dp = [INF] * (n + 1)
    dp[start] = 0
    heap = [(0, start)]

    while heap:
        weight, current = heappop(heap)

        if dp[current] < weight:
            continue

        for w, next in graph[current]:
            next_weight = weight + w

            if next_weight < dp[next]:
                dp[next] = next_weight
                heappush(heap, (next_weight, next))

    return dp

for _ in range(int(input())):
    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((d, b))
        graph[b].append((d, a))

    candidate = [int(input()) for _ in range(t)]

    start_dp = dijkstra(s)
    g_dp = dijkstra(g)
    h_dp = dijkstra(h)

    answer = []
    
    for i in candidate:
        path1 = start_dp[g] + g_dp[h] + h_dp[i]
        path2 = start_dp[h] + h_dp[g] + g_dp[i]
        
        if path1 == start_dp[i] or path2 == start_dp[i]:
            answer.append(i)

    answer.sort()

    print(*answer)

INF값을 float(‘inf’)로 하면 틀립니다.