[Python] 1854. k번째 최단거리 구하기(플4) / 차근차근 예제 뜯기

목차

1. 문제

2. 핵심 아이디어

3. 코드

4. 배운 점

---

 

1. 문제

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

1854번: K번째 최단경로 찾기

---

2. 핵심 아이디어

 

2-1) distance보다 긴 거리값이 들어와도 skip X

 

기존의 다익스트라는 우선순위 큐에서 나온 거리가

distance에 이미 갱신되어 있는 거리보다 멀 경우 continue를 통해 연산을 스킵했다.

 

<기존 코드>

 

그러나 이번 문제는 조금 다르다.

 

이 노드에 도달하는 거리가 기존에 갱신 값보다 더 길어도

그 기록을 계속 가지고 있어야 한다.

 

그렇게 노드에 도달하는 상위 k개의 기록을 계속 가지고 있어야 하므로

distance 를 2차원 배열로 만들어 주어야 한다.

---

2-2) 구체적인 갱신과정

 

만약 어떤 노드에 이를 수 있는 거리가 3가지로 나뉘어져있고

3번째로 최단거리를 구하는 공식을 생각해보자

 

노드 X에 이를 수 있는 거리의 경우 = [3,7,8]

 

여기서 x를 6에 이를 수 있는 새로운 루트를 발견했다면

나는 무엇을 비교해야 할까?

 

먼저 상위 3위 안에 속하는지 확인하기 위해

루트 거리 중 3등의 거리(8)과 비교해볼 것이다.

 

그럼 8보다 6이 더 작으므로 3등 안에 속한다는 결론이 났으므로

8대신 6을 대체해준다.

 

노드 X에 이를 수 있는 거리의 경우 = [3,7,8]

=> 노드 X에 이를 수 있는 거리의 경우 = [3,7,6]

 

이제 이 거리 배열을 다시 정렬해주면 순서대로 정렬된다.

 

노드 X에 이를 수 있는 거리의 경우 = [3,6,7]

 

이런식의 검토를 다익스트라를 통해 매 노드 마다 해주면 된다.

 즉 우리가 검토해야 하는 질문은 다음과 같다.

 

->상위 k 개 안에 속하는 가?

-> 속하면 꼴등을 대체해주고 정렬

-> 안속하면 skip

 

---

2-3) 예시를 통한 갱신과정

 

주어진 예시는 다음과 같다.

먼저 distance를 다음과 같이 초기화해준다.

노드번호	1번째 최단거리	2번째 최단거리
1	0	INF
2	INF	INF
3	INF	INF
4	INF	INF
5	INF	INF

=> 1번과 연결된 노드들을 대상으로 dijkstra를 실행하면 다음과 같이 최단거리가 갱신된다.

 

노드번호	1번째 최단거리	2번째 최단거리
1	0	INF
2	2	INF
3	7	INF
4	5	INF
5	6	INF

 

여기서 4번 노드를 중심으로 보자

dijkstra 과정에서 4번노드에는 다음 루트의 케이스가 계산된다.

노드번호	1번째 최단거리	2번째 최단거리
1	0	INF
2	2	INF
3	7	INF
4	5	INF
5	6	INF

=> 그럼 기존에 4 노드의 거리 리스트를 바탕으로 13이라는 거리가  상위 2위 안에 속하는지 봐야 한다.

=> 4노드의 2등 거리가 INF이므로 상위 2등에 속한다.

 

그럼 2등 거리를 13으로 갱신하고 정렬해주면 다음과 같이 된다.

노드번호	1번째 최단거리	2번째 최단거리
1	0	INF
2	2	INF
3	7	INF
4	5	13
5	6	INF

 

 

계속해서 dijkstra를 하다보면 다음 루트도 나온다.

=> 이경우에는 거리가 4이다.

 

노드번호	1번째 최단거리	2번째 최단거리
1	0	INF
2	2	INF
3	7	INF
4	5	13
5	6	INF

=> 기존에 4 노드의 거리 리스트를 바탕으로 4라는 거리가  상위 2위 안에 속하는지 봐야 한다.

=> 4노드의 2등 거리가 13이므로 상위 2등에 속한다.

 

그럼 2등 거리를 4로 갱신하고 정렬해주면 다음과 같이 된다.

노드번호	1번째 최단거리	2번째 최단거리
1	0	INF
2	2	INF
3	7	INF
4	4	5
5	6	INF

 

 

이런 과정을 반복한 후 2번째 최단거리에 대해서만 출력해주면 된다.

---

3. 코드

#1854. k번째 최단경로 찾기(플4)
#distance를 2차원으로 만들어 보기

import heapq
import sys
#input= lambda: sys.stdin.readline().strip('\n')
#최대값으로 초기화
INF= sys.maxsize

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

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

#간선 입력받기
graph=[[] for _ in range(n+1)]

for _ in range(m):
  a,b,c= map(int, input().split())
  graph[a].append((b,c))

def dijkstra(start):
  q=[]
  distance[start][0]=0
  heapq.heappush(q, (0,start))

  while(q):
    dist, now= heapq.heappop(q)

    for next_node, next_cost in graph[now]:
      cost= next_cost+dist
	  
      #상위 k번째 안에 속하는가?
      if cost<distance[next_node][-1]:
      		
            #속하면 꼴등이랑 switch ->정렬
            distance[next_node][-1]=cost
            distance[next_node].sort()
            
            heapq.heappush(q, (cost, next_node))

dijkstra(1)

#k번째만 출력해주기
for i in distance[1:]:
  if i[k-1]>=INF:
    print(-1)
  else:
    print(i[k-1])

---

4.배운 점

그래프 문제의 난이도를 조금 올리다보니 dijkstra를 응용하여 복잡한 조건 안에서 최단거리를 갱신하는 식의 문제가 자주 보이는 것 같다. 알고리즘을 구현하는 것보다 알고리즘이 담고 있는 의미를 정확하게 파악해야 응용이 가능한 것 같다.

 

위의 문제도 distance를 갱신하기 위해선 heap에서 push와 pop을 통해 나오는 값들의 의미, 어떤 부분을 비교하고 있고, 왜 비교하는지를 정확히 알아야 하는 문제였다.

 

특히, 2차원 행렬 문제가 나오면 플로이드와 다익스트라 중에 고민하게 되는데 플로이드는 시간복잡도가 좋지 않으니 다익스트라를 우선적으로 고려하는 습관을 길러야 겠다는 생각이 들었다.