이번 문제는 다익스트라를 통해 풀었다.

 

원래 처음에는 경로를 추적할 수 있도록, 거리를 측정하는 배열에 [이전 노드, 거리] 값을 저장하도록 했었다.

문제는 다익스트라는 해당 노드까지 최단거리로 갈 수 있는 경로가 많은 경우, g와 h를 포함하지 않을 수 있기 때문에 문제가 되었다.

 

그래서 s -> g -> h -> 목적지, s -> h -> g -> 목적지를 구하기 위해, s ,g, h를 시작점으로 각각 다익스트라 통해 최단 거리를 구한다.

 

import heapq
import sys

from collections import deque



def djk(board, s, n):

    hq = []
    
    dst = deque([987654321 for _ in range(n+1)])
    
    heapq.heappush(hq, (0, s))
    dst[s] = 0
    
    
    while hq:
        d, node = heapq.heappop(hq)
           
        if dst[node] < d:
            continue
        
        for i in board[node]:
            
            if board[node][i] + d < dst[i]:
                dst[i] = board[node][i] + d
                heapq.heappush(hq, (dst[i], i))
    
    return dst


T = int(sys.stdin.readline())


for _ in range(T):
    n, m , t = map(int, sys.stdin.readline().split())
    s, g, h = map(int, sys.stdin.readline().split())
    
    board = [{} for _ in range(n+1)]
    
    for _ in range(m):
        a, b, d = map(int, sys.stdin.readline().split())
        
        board[a][b] = d
        board[b][a] = d
    
    dst_condidate = deque()
    
    for _ in range(t):
        dst_condidate.append(int(sys.stdin.readline()))
    
    ss = djk(board, s, n)
    gg = djk(board, g, n)
    hh = djk(board, h, n)
    
    ans = []
    
    for dd in dst_condidate: 
        if ss[dd] == ss[g] + board[g][h] + hh[dd] or ss[dd] == ss[h] + board[h][g]+ gg[dd]:
            ans.append(dd)
            
    ans.sort()
    for w in ans:
        print(w, end=' ')
    print()

원래 ans에 str(dd)를 append 하고 join을 이용하려 했었는데, str은 sort가 integer 형이랑 결과가 다르게 나온다. 이 점에서 매우 시간을 많이 소비했다.

반응형

+ Recent posts