[백준] 9370번 미확인 도착지 파이썬 코드

코드

import heapq, sys

input = sys.stdin.readline
INF = 987654321

# 다익스트라 알고리즘
def dijkstra(x,graph):
  visited = [INF]*(n+1)
  queue = [(0,x)]
  visited[x] = 0
  while queue:
    cost, now = heapq.heappop(queue)
    if cost > visited[now]: continue
    for i in graph[now]:
      newcost = cost + i[1]
      if visited[i[0]] > newcost:
        visited[i[0]] = newcost
        heapq.heappush(queue, (newcost, i[0]))
  return visited

tc = int(input())

for _ in range(tc):
  n,m,t = map(int, input().split())
  s,g,h = map(int, input().split())
  arrive = [] # 도착 후보지
  res = [] # 정답
  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))
  for _ in range(t):
    arrive.append(int(input()))
    
  Ds = dijkstra(s,graph) # s에서 출발한 최단거리
  Dg = dijkstra(g,graph) # g에서 출발한 최단거리
  Dh = dijkstra(h,graph) # h에서 출발한 최단거리
  for x in arrive:
    if Ds[g]+Dg[h]+Dh[x] == Ds[x] or Ds[h]+Dh[g]+Dg[x]==Ds[x]:
    # s -> g -> h -> x로 이동 or s -> h -> g -> x로 이동
      res.append(x)
  res.sort() # 오름차순 정렬
  print(*res)

처음 접근은 도착지에서부터 출발지로 가면서 지나야 하는 길을 거치면 flag를 True로 바꿔주는 식이었다. 그런데 그 길을 지나도 최단거리가 다른 길에서 초기화될 수 있기 때문에, 구멍이 있다는 것을 발견했다. 어떻게 g-h를 지나는지 고민했는데, 잘 모르겠어서 정답을 참고했다... g-h를 확정적으로 지나게 하는 방식으로 짜는 방법이 있었다. 왜 이걸 생각 못했을까 싶기도 하지만, 부족해서 그런 것이니 더 열심히 공부 해야겠다고 느꼈다.