[백준] 1504 특정한 최단경로 python

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

1504번: 특정한 최단 경로

지난번 풀었던 문제인데 재채점된 이후에 값이 틀리게 나와서 다시 풀게되었다.

---

import sys; input = sys.stdin.readline
from heapq import heappop, heappush

def dijkstra(start, end, num_of_nodes):
    INF = 200_000 * 1_000 * 2
    q = []
    q.append((0, start))
    visit = [INF for _ in range(num_of_nodes + 1)]
    visit[start] = 0
    while q:
        dist, cur = heappop(q)
        for adj, adj_dist in linked_list[cur]:
            tmp_dist = dist + adj_dist
            if visit[adj] > tmp_dist:
                visit[adj] = tmp_dist
                heappush(q, (tmp_dist, adj))
    # print(visit)
    if visit[end] < INF:
        return visit[end]
    return None


# 0 입력
N, E = map(int, input().split())
linked_list = [list() for _ in range(N+1)]
for _ in range(E):
    a, b, c = map(int, input().split())
    linked_list[a].append((b, c))
    linked_list[b].append((a, c))
v1, v2 = map(int, input().split())

# 1 최단경로 탐색
answer = -1

d1, d2, d3 = dijkstra(1, v1, N), dijkstra(v1, v2, N), dijkstra(v2, N, N)
d4, d5, d6 = dijkstra(1, v2, N), dijkstra(v2, v1, N), dijkstra(v1, N, N)

ans1, ans2 = None, None
if d1 and d2 and d3:
    ans1 = d1 + d2 + d3
if d4 and d5 and d6:
    ans2 = d4 + d5 + d6

if ans1 and ans2:
    answer = min(ans1, ans2)
elif ans1:
    answer = ans1
elif ans2:
    answer = ans2

print(answer)

이렇게 하면 지난번 풀었던 문제와 같은 풀이일것이다.

하지만 80 퍼센트쯤 틀린 답이 나오게 되는데 그 이유는 INF 와 근접한 값이 있으면 answer 값이 나올수도 있다.

 

 

---

import sys; input = sys.stdin.readline
from heapq import heappop, heappush
INF = 200_000 * 1_000 * 2

def dijkstra(start, end, num_of_nodes):
    q = [(0, start)]
    visit = [INF for _ in range(num_of_nodes + 1)]
    visit[start] = 0
    while q:
        dist, cur = heappop(q)
        if visit[cur] < dist:
            continue
        for adj, adj_dist in linked_list[cur]:
            if visit[adj] > dist + adj_dist:
                visit[adj] = dist + adj_dist
                heappush(q, (visit[adj], adj))
    return visit[end]


# 0 입력
N, E = map(int, input().split())
linked_list = [list() for _ in range(N+1)]
for _ in range(E):
    a, b, c = map(int, input().split())
    linked_list[a].append((b, c))
    linked_list[b].append((a, c))
v1, v2 = map(int, input().split())

# 1 최단경로 탐색
ans1 = dijkstra(1, v1, N) + dijkstra(v1, v2, N) + dijkstra(v2, N, N)
ans2 = dijkstra(1, v2, N) + dijkstra(v2, v1, N) + dijkstra(v1, N, N)
answer = min(ans1, ans2)
if answer >= INF:
    answer = -1
print(answer)

이렇게 다익스트라 안에서 판단을 하지 않고 바깥에서 INF 보다 작은 경우만 값을 내놓는다면 성공한다.