BOJ-1854 - Python

제일 빠른 경로만을 찾던 다익스트라 문제들과는 달리 k번째로 빠른 경로를 찾는 문제. 첫 번째 시도는 n개의 heap을 이용하여 정렬하며 해당 배열의 크기가 k보다 작을 때만 q에 삽입하는 형식으로 풀었으나 틀렸습니다 결과를 받게 되었다. 내가 생각하는 반례로는 다른 지점들이 k 이상이 되어야 k번째 방문 이력이 생기는 경우이다. 두 번째 풀이는 k개의 INF로 이루어진 배열이 n+1개인 2차원 배열을 이용하였다. k-1번째 수와 비교하여 비용이 작을 때 q에 삽입한 후, k-1번째 수에 삽입, 해당 배열을 정렬하여 풀었다. heap을 이용하면 더욱 빠르게 풀 수 있을 것 같으나 INF로 초기화하는 과정때문에 오히려 연산이 늘어날 것만도 같다.

import heapq
import sys
INF = sys.maxsize
input = sys.stdin.readline

def solve(n, k, graph):

    q = []
    heapq.heappush(q, [0, 1])
    # k개의 INF로 이루어진 배열이 n+1개
    dists = [[INF for _ in range(k)] for _ in range(n+1)]
    dists[1][0] = 0

    while q:
        cnt, now = heapq.heappop(q)
        for np, tc in graph[now]:
            nc = tc + cnt

            # k번째 수와 비교
            if dists[np][k-1] > nc:
                heapq.heappush(q, [nc, np])
                # k번째 수에 삽입 및 정렬
                dists[np][k-1] = nc
                dists[np].sort()

    return dists

n, m, k = map(int, input().split())
graph = [[] for _ in range(n+1)]

for _ in range(m):
    s, e, w = map(int, input().split())
    graph[s].append([e, w])

dists = solve(n, k, graph)

for i in range(1, n+1):
    print(-1 if dists[i][k-1] == INF else dists[i][k-1])

Written by@hyesungoh
Learning every moment

InstagramGitHubLinkedIn