[Python] boj no. 9370 : 미확인 도착지

Problem

https://www.acmicpc.net/problem/9370

Idea

도착지가 주어질 때 출발점으로부터 각각의 목적지로의 최단 경로가 정점 g와 h 사이의 간선을 포함하는지를 알아내는 문제이다.

a로부터 b까지의 최소 비용을 dist(a, b)라고 할 때, 예상 목적지 x가 출발점 s로부터 정점 g와 h사이의 간선을 포함하는 최단 경로를 갖는 경우 다음과 같은 조건 중 하나가 만족해야 한다.

1. dist(s, x) == dist(s, g) + dist(g, h) + dist(h, x)
2. dist(s, x) == dist(s, h) + dist(h, g) + dist(g, x)

따라서 정점 s, g, h에 대해 다익스트라 알고리즘을 수행하고 모든 예상 목적지에 대해 조건을 비교하여 풀이했다.

Code

import sys
import heapq

def dijkstra(start):
    dist = [sys.maxsize for _ in range(n+1)]
    dist[start] = 0
    heap = []
    heapq.heappush(heap, (dist[start], start))
    
    while heap:
        pd, pn = heapq.heappop(heap)
        
        if pd > dist[pn]:
            continue
        
        for nd, nn in edge[pn]:
            if dist[pn] + nd < dist[nn]:
                dist[nn] = dist[pn] + nd
                heapq.heappush(heap, (dist[nn], nn))
                
    return dist
    
input = sys.stdin.readline

T = int(input())
for _ in range(T):
    n, m, t = map(int, input().split())
    s, g, h = map(int, input().split())
    edge = [[] for _ in range(n+1)]
    for _ in range(m):
        a, b, d = map(int, input().split())
        edge[a].append((d, b))
        edge[b].append((d, a))
    
    xlist = []
    for _ in range(t):
        xlist.append(int(input()))
        
    sol = []
    s_dist = dijkstra(s)
    g_dist = dijkstra(g)
    h_dist = dijkstra(h)
    
    for x in xlist:
        if g_dist[h] + s_dist[g] + h_dist[x] == s_dist[x] or g_dist[h] + s_dist[h] + g_dist[x] == s_dist[x]:
            sol.append(x)
        
    sol.sort()
    print(*sol)

Comment