이것이 코딩테스트다 with 파이썬 - Chp9. 최단 경로_2. 미래 도시

Try 1

import heapq

INF = int(1e9)
n, m = map(int, input().split())
graph = []
for i in range(m):
    a,b = (map(int, input().split()))
    graph.append((a,b))
    # graph.append(list(map(int, input().split())))
x, k = map(int, input().split())
distance = [INF] * (n+1)

def dijkstra(start, end):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + 1
            if cost < distance[i]:
                distance[i] = cost
                heapq.heappush(q, (cost, i))
    return distance[end]

result = dijkstra(start=1 ,end=4) + dijkstra(start=4, end=5)
print(result)

5 7
1 2
1 3
1 4
2 4
3 4
3 5
4 5
4 5
-> 1000000002

==> 정답은 3이 나와야 한다. 아마 result에서 start와 end 지점에 문제가 생긴 것으로 보인다.

Try 2.

import heapq

INF = int(1e9)
n, m = map(int, input().split())
graph = []
for i in range(m):
    a,b = (map(int, input().split()))
    graph.append((a,b))
    # graph.append(list(map(int, input().split())))
x, k = map(int, input().split())
distance = [INF] * (n+1)

def dijkstra(start, end):
    q = []
    heapq.heappush(q, (0, start))
    distance[start] = 0
    while q:
        dist, now = heapq.heappop(q)
        if distance[now] < dist:
            continue

        for i in graph[now]:
            cost = dist + 1
            if cost < distance[i]:
                distance[i] = cost
                heapq.heappush(q, (cost, i))

    if distance[end] != INF: 
        return distance[end]
    else:
        return -1 
result = dijkstra(start=0 ,end=(x-1)) + dijkstra(start=x-1, end=k-1)

print(result)

첫번째 입력 예시에선 정답 '3'을 도출하였다.
두번째 입력 예시인
4 2
1 3
2 4
3 4
->line 22, in dijkstra
for i in graph[now]:
IndexError: list index out of range

정답 예시

import sys
input = sys.stdin.readline
INF = int(1e9)

n, m = map(int, input().split())

graph = [[INF] * (n+1) for _ in range(n+1)]

for i in range(1, n+1):
	for j in range(1, n+1):
		if i == j:
			graph[i][j] = 0

for i in range(m):
	a,b = map(int,input().split())
	graph[a][b] = 1
    graph[b][a] = 1

x, k = map(int, input().split())

for k in range(1, n+1):
	for i in range(1, n+1):
		for j in range(1, n+1):
			graph[i][j] = min(graph[i][j], graph[i][k]+graph[k][j])

distance = graph[1][k] + graph[k][x]

if distance >= INF :
	print(-1)
else :
	print(distance)