[python]9370 미확인 도착지, 다익스트라 반복 돌리기

전에 풀었던 특정한 최단경로 문제를 응용한 문제였다.

https://weallbinna.tistory.com/manage/newpost/172?type=post&returnURL=https%3A%2F%2Fweallbinna.tistory.com%2Fmanage%2Fposts%2F 

https://weallbinna.tistory.com/manage/newpost/172?returnURL=https%3A%2F%2Fweallbinna.tistory.com%2Fmanage%2Fposts%2F&type=post

문제에서 출발점부터 목적지까지의 최단경로와 반드시 거쳐야하는 길을 지나면서 최단경로가 같은지를 확인하는 문제였다.

 

그래서 2번의 특정한 최단경로를 확인하기전에 후보 목적지들의 최단경로를 구해주면 문제는 풀리게 되었다.

 

from heapq import heappush,heappop
import sys
input = sys.stdin.readline
def dix(start,end):                       #다익스트라 돌리기
    que = [(0,start)]
    visit = [1e9 for i in range(node+1)]
    visit[start] = 0
    while que:
        curlen,curx = heappop(que)
        if curx == end:    #목적지에 다다를 수 있다면
            return (curlen,curx)
        for nextx,nextlen in graph[curx]:
            if visit[nextx]>visit[curx]+nextlen:
                visit[nextx] = visit[curx]+nextlen
                heappush(que,(visit[nextx],nextx))
    return (1e9,1)  #목적지에 다다를 수 없다면
    
for _ in range(int(input().strip())):
    node,edge,dest = map(int,input().split())
    start,ess1,ess2 = map(int,input().split())
    graph = [[]for i in range(node+1)]
    destl = []
    for i in range(edge):
        fst,scd,trd = map(int,input().split())
        graph[fst].append((scd,trd))
        graph[scd].append((fst,trd))
    for i in range(dest):
        destl.append(int(input().strip()))
    std = [0 for i in range(node+1)]
    for i in destl:
        std[i] = dix(start,i)[0]    #후보 목적지들 최단경로 구해주기
    attp = []
    for i in destl:   #실질적으로 특정한 길을 지나면서도 최단경로가 유지되는지 확인하는 작업
        tmp = min(dix(start,ess1)[0]+dix(ess1,ess2)[0]+dix(ess2,i)[0],
                 dix(start,ess2)[0]+dix(ess2,ess1)[0]+dix(ess1,i)[0])
        if tmp==std[i]:   #만약 맞다면 결과값에 갱신
            attp.append(i)
    print(*sorted(attp))

그러나 내 풀이의 문제는 다익스트라를 너무 자주 돌려서 시간이 오래걸린다.

내가 최적화 할려고 리스트를 반환할려 했지만 메모리가 초과되서 어려움을 겪었다.

다른사람들 풀이 보니 3번만 돌려서 리스트에서 g,h를 인덱싱해서 푸는 방법도 있었다. 

gptpem38님의 풀이

import sys
import heapq

input = sys.stdin.readline

T = int(input())

def run(start):
    Q = []
    distances = [9876543210] * (n + 1)
    distances[start] = 0
    heapq.heappush(Q, (0, start))

    while Q:
        cur_dist, cur_node = heapq.heappop(Q)
        if cur_dist > distances[cur_node]:
            continue
        for next_node, next_dist in G[cur_node]:
            sum_dist = cur_dist + next_dist
            if sum_dist < distances[next_node]:
                distances[next_node] = sum_dist
                heapq.heappush(Q, (sum_dist, next_node))
    return distances

for _ in range(T):
    n, m, t = map(int, input().strip().split())
    s, g, h = map(int, input().strip().split())

    G = {k: [] for k in range(1, n+1)}
    for _ in range(m):
        # from, to, distance
        a, b, d = map(int, input().strip().split())
        G[a].append((b, d)) # heap ? or d, b instead ?
        G[b].append((a, d))
    candidates = [int(input()) for _ in range(t)]   #목적지 후보들

    s_distances = run(s)   #s부터 다익스트라
    g_distances = run(g)   #g부터
    h_distances = run(h)   #h부터
    ans = []
    for candidate in candidates:
        if s_distances[g] + g_distances[h] + h_distances[candidate] == s_distances[candidate]:
            ans.append(candidate)  # g-h-목적지가 해당하는지
        elif s_distances[h] + h_distances[g] + g_distances[candidate] == s_distances[candidate]:
            ans.append(candidate)  # h-g-목적지가 해당하는지
    print(*sorted(ans))