백준 9370 - 미확인 도착지 (Python)

https://acmicpc.net/problem/9370 

9370번: 미확인 도착지

일반 다익스트라에 더해 내가 해당 지점에 방문했었는지 여부를 검사한다.
단순하게 일반적인 다익스트라에 도착 여부를 더해주면 된다.

import sys
from heapq import heappop, heappush

input=sys.stdin.readline
INF=sys.maxsize

T=int(input())
for _ in range(T):
    answers=[]
    n,m,t=map(int,input().strip().split())
    s,g,h=map(int,input().strip().split())
    graph=[[] for _ in range(n)]
    for _ in range(m):
        a,b,d=map(int,input().strip().split())
        graph[a-1].append([d,b-1])
        graph[b-1].append([d,a-1])
    arrivals=set()
    for _ in range(t):
        arrivals.add(int(input()))
    dp=[(INF,False) for _ in range(n)]
    q=[(0,s-1,False)]
    while q:
        c_v,c_n,c_p=heappop(q)
        for next_vn in graph[c_n]:
            n_v,n_n=next_vn
            temp = c_p
            if (c_n == g - 1 and n_n == h - 1) or (n_n == g - 1 and c_n == h - 1):
                temp = True
            if  dp[n_n][0] > c_v+n_v:
                dp[n_n]=(c_v+n_v,temp)
                heappush(q,(c_v+n_v,n_n,temp))
            if not dp[n_n][1] and temp and c_v+n_v == dp[n_n][0]:
                dp[n_n] = [c_v + n_v, temp]
                heappush(q,(c_v+n_v,n_n,temp))

    for i in arrivals:
        if dp[i-1][1]:
            answers.append(i)

    answers.sort()
    print(*answers)
        print(i, end=' ')