[CLASS 4]백준 1504번 - 특정한 최단 경로

1504번 - 특정한 최단 경로

 

다익스트라를 사용하는 문제이다. 그런데 경유지가 하나면 모르겠는데 2개여서 무조건 다익스트라를 여러번 사용해야 한다.

 

내 코드:

# dawitblog.tistory.com/150
from heapq import heappop,heappush
from sys import stdin
input = stdin.readline
INF = 1e9

n, m = map(int,input().split())
nodes = [[] for i in range(n+1)]
for _ in range(m):
    a, b, d = map(int,input().split())
    nodes[a].append((b,d))
    nodes[b].append((a,d))
# 경유지
s1, s2 = map(int,input().split())

# 간선 없을 시
if m == 0:
    print(-1)
    exit()

# 다익스트라. end 지점까지의 최단거리 확정시 바로 종료
def find(start,end,nodes):
    q = []
    heappush(q,(0,start))
    distance = [INF]*(n+1)
    distance[start] = 0
    while True:
        dist, now = heappop(q)
        if distance[now] < dist:
            continue
        for i in nodes[now]:
            cost = dist + i[1]
            if cost < distance[i[0]]:
                distance[i[0]] = cost
                heappush(q,(cost,i[0]))
        if now == end:
            break
    return distance[end]

# s1과 s2까지의 최단경로를 각각 저장
q = []
heappush(q,(0,1))
distance = [INF]*(n+1)
distance[1] = 0
count = 0
while True:
    dist, now = heappop(q)
    if distance[now] < dist:
        continue
    if now == s1:
        count += 1
    if now == s2:
        count += 1
    for i in nodes[now]:
        cost = dist + i[1]
        if cost < distance[i[0]]:
            distance[i[0]] = cost
            heappush(q,(cost,i[0]))
    if count == 2:
        break

# d1 = 1->s1 , d2 = 1->s2
d1 = distance[s1]
d2 = distance[s2]

# b = s1->s2 , a1 = s1->n , a2 = s2->n
b = find(s1,s2,nodes)
a1 = find(s1,n,nodes)
a2 = find(s2,n,nodes)

# 출력
ans = min(d1+b+a2,d2+b+a1)
print(ans) if ans != 1e9 else print(-1)

 

나는 목적지까지의 최단경로를 구했다면 바로 탐색을 종료하는 방식의 다익스트라를 총 4번 사용했다.

 

s1과 s2사이의 거리 b를 이런 식으로 따로 구하지 않고 그냥 완전한 다익스트라 3번이면 문제를 풀 수 있긴 하지만 이 문제에 한해서는 큰 차이를 내는것 같지는 않다.