[백준] 9370번-(Python 파이썬) - Dijkstra

문제링크 : https://www.acmicpc.net/problem/9370

처음엔 문제에 입력값이 너무 많아 이해하는데 시간이 조금 걸렸었다.
문제를 처음부터 찬찬히 읽고 다 이해한 뒤에 s에서 g, h를 거쳐서 가는 최소값과 s에서 바로 가는 최소값이 같으면 정답이 된다는 걸 나중에 도출할 수 있었다.
g를 먼저 지나갈 때와 h를 먼저 지나갈 때 둘 다 구해줘야 한다.
조건에 맞게 정렬 후 출력을 해주면 된다.

import sys
from heapq import heappush, heappop

input = sys.stdin.readline
inf = sys.maxsize

def dijkstra(start):
    heap = []
    heappush(heap, [0, start])
    dp = [inf] * (n + 1)
    dp[start] = 0
    while heap:
        w, now = heappop(heap)
        for n_n, wei in graph[now]:
            n_w = wei + w
            if n_w < dp[n_n]:
                dp[n_n] = n_w
                heappush(heap, [n_w, n_n])
    return dp

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])

    x = []
    for _ in range(t):
        x.append(int(input()))

    s_ = dijkstra(s)
    g_ = dijkstra(g)
    h_ = dijkstra(h)
    ans = []

    for i in x:
        if s_[g] + g_[h] + h_[i] == s_[i] or s_[h] + h_[g] + g_[i] == s_[i]:
            ans.append(i)
    
    ans.sort()
    print(*ans)

좋은 웹페이지 즐겨찾기