파이썬 코딩테스트 강좌, K번째 최단 경로 찾기

알고리즘 문제 해결에 대한 이해와 연습은 소프트웨어 개발자, 특히 취업을 목표로 하는 학생들에게 필수적입니다. 다양한 유형의 문제를 풀어보면서 문제 해결 능력을 키우고 알고리즘과 자료 구조를 실질적으로 사용하는 법을 배워야 합니다. 이번 강좌에서는 “K번째 최단 경로 찾기”라는 주제를 다루고, 해당 문제를 파이썬으로 해결하는 과정을 상세히 설명하겠습니다.

문제 설명

그래프가 주어졌을 때, 특정 두 노드 사이의 K번째 최단 경로를 찾는 문제입니다. K번째 최단 경로란, 두 노드 사이의 최단 경로 중에서 K번째로 짧은 경로를 의미합니다. 이 문제는 Dijkstra 알고리즘과 같은 최단 경로 탐색 알고리즘의 변형을 통해 접근할 수 있습니다.

문제 요약

그래프는 비가중치 및 가중치를 가진 방향성 그래프입니다.
Graph가 주어짐 (노드 수 N, 간선 수 M).
시작 노드 S와 끝 노드 E가 주어진다.
K번째 최단 경로를 구해야 한다.

입력 형식

N M K
a1 b1 c1
a2 b2 c2
...

위의 입력 형식에서, N은 노드의 수, M은 간선의 수, K는 찾고자 하는 K번째 최단 경로를 의미합니다. 각 간선은 세 개의 숫자로 주어지며, a는 시작 노드, b는 종료 노드, c는 가중치를 뜻합니다.

출력 형식

K번째 최단 경로의 길이 또는 K번째 최단 경로가 존재하지 않으면 -1

예제

입력 예제

출력 예제

풀이 과정

이 문제를 해결하기 위해서 K번째 최단 경로를 찾는 알고리즘을 구현해야 합니다. 일반적인 최단 경로 알고리즘인 Dijkstra 알고리즘을 변형하여 사용할 예정이며, 각 경로의 비용을 관리하는 방법을 고안해야 합니다. 여기서는 우선순위 큐(Priority Queue)를 활용하여 경로를 탐색하는 방식으로 해결하겠습니다.

단계별 접근

1단계: 그래프 표현하기

그래프를 인접 리스트 형태로 표현합니다. 각 노드와 간선 정보를 담기 위해 딕셔너리나 리스트를 사용합니다.

from collections import defaultdict
import heapq

def build_graph(edges):
    graph = defaultdict(list)
    for (u, v, w) in edges:
        graph[u].append((v, w))
    return graph

2단계: 우선순위 큐를 사용한 경로 탐색

우선순위 큐를 사용하여 최단 경로를 탐색합니다. 각 경로의 비용과 노드를 튜플 형태로 저장하여 관리할 수 있습니다. 이때, K개의 경로를 저장하기 위한 리스트를 준비해야 합니다.

def kth_shortest_path(graph, start, end, k):
    # 초기화
    pq = []
    heapq.heappush(pq, (0, start))  # (비용, 노드)
    paths = defaultdict(list)

    while pq:
        cost, node = heapq.heappop(pq)
        
        # K번째 경로를 찾았을 경우
        if node == end:
            paths[end].append(cost)
            if len(paths[end]) >= k:
                break
        
        for neighbor, weight in graph[node]:
            new_cost = cost + weight
            heapq.heappush(pq, (new_cost, neighbor))
    
    return paths[end][k - 1] if len(paths[end]) >= k else -1

3단계: 전체 알고리즘 통합하기

위의 두 기능을 결합하여 전체 알고리즘을 완성합니다. 입력을 받고 그래프를 만들고 K번째 최단 경로를 찾는 로직을 구현하겠습니다.

def main():
    N, M, K = map(int, input().split())
    edges = [tuple(map(int, input().split())) for _ in range(M)]
    graph = build_graph(edges)
    result = kth_shortest_path(graph, 1, N, K)
    print(result)

if __name__ == "__main__":
    main()

결론 및 마무리

이번 강좌에서는 K번째 최단 경로 찾기 문제를 다루어 보았습니다. 그래프의 구조와 Dijkstra 알고리즘의 변형 방법을 이해하고, 우선순위 큐를 통해 경로를 탐색하는 방법을 익혔습니다. 알고리즘 문제는 다양하게 변형되어 출제될 수 있으므로, 여러 문제를 많이 풀어보는 것이 중요합니다.

이제 여러분은 주어진 문제를 해결할 수 있는 기반을 다졌습니다. K번째 최단 경로 문제뿐만 아니라 다양한 그래프 관련 문제에서도 활용할 수 있는 방법들이므로, 응용력을 키우는 것이 중요합니다. 앞으로도 더 많은 알고리즘과 문제 해결 방법을 연습하시길 바랍니다.

Happy Coding!