[코딩 테스트 연습(파이썬/Python)] 백준(BOJ) 13907번 _ 세금

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

13907번: 세금

[문제]

주언이는 경제학을 배워 행상인이 되었다. 두 도시를 오가며 장사를 하는데, 통행료의 합이 가장 적은 경로로 이동하려 한다. 도시들은 양방향 도로로 연결되어있으며, 도로마다 통행료가 존재한다.

그런데 정부는 세금 인상안을 발표하였다. 세금을 한 번에 올리면 문제가 발생할 수 있으므로 여러 단계에 걸쳐서 올린다고 한다. 세금이 A만큼 오르면 모든 도로의 통행료가 각각 A만큼 오르게 된다. 세금이 오르게 되면 주언이가 내야 하는 통행료가 변할 수 있다.

주언이를 도와 초기의 최소 통행료와 세금이 오를 때마다의 최소 통행료를 구하시오.

 

[입력]

첫 번째 줄에 세 정수 N (2 ≤ N ≤ 1,000), M (1 ≤ M ≤ 30,000), K (0 ≤ K ≤ 30,000)가 주어진다. 각각 도시의 수, 도로의 수, 세금 인상 횟수를 의미한다.

두 번째 줄에는 두 정수 S와 D (1 ≤ S, D ≤ N, S ≠ D)가 주어진다. 각각 출발 도시와 도착 도시 번호를 의미한다. 도시 번호는 1부터 시작한다.

다음 M개의 줄에는 각각 도로 정보를 나타내는 세 정수 a, b (1 ≤ a < b ≤ N), w (1 ≤ w ≤ 1,000)가 주어진다. 도시 a와 도시 b가 통행료 w인 도로로 연결되어 있다는 것을 의미한다.

다음 총 K개의 줄에는 각각 정수 p (1 ≤ p ≤ 10)가 주어진다. 각각 첫 번째, 두 번째, …, K 번째에 인상되는 세금을 의미한다.

S에서 D로 이동할 수 없는 경우는 주어지지 않는다.

[출력]

첫 번째 줄에 세금이 오르기 전의 최소 통행료를 출력한다.

두 번째 줄부터 K개의 줄에 각각 첫 번째, 두 번째, …, K 번째 세금이 올랐을 때의 최소 통행료를 출력한다.

 

[아이디어]

1) 다익스트라를 이용해 최단 거리를 구하되, 사용하는 경로의 수(1개 ~ n-1개)에 따른 모든 최소 비용을 구한다.

2) 해당 최소의 비용을 기준으로 세금 인상에 따른 필요 비용을 업데이트하며 최솟값을 출력한다.

 

import heapq
import sys
input=sys.stdin.readline
n,m,k=map(int,input().split())
s,e=map(int,input().split())
g=[[] for _ in range(n+1)]
for _ in range(m):
    a,b,w=map(int,input().split())
    g[a].append((b,w))
    g[b].append((a,w))
INF=int(1e9)
dist=[[INF]*(n+1) for _ in range(n+1)]
def dijk(s,e,dist):
    q = []
    heapq.heappush(q,(0,s,0))
    dist[s] = [0]*(n+1)
    while q:
        d, now, cnt = heapq.heappop(q)
        if dist[now][cnt] < d:
            continue
        for x,w in g[now]:
            cost = d + w
            if cost < dist[x][cnt+1]:
                dist[x][cnt+1] = cost
                heapq.heappush(q,(cost,x,cnt+1))
    return dist[e]

a = dijk(s,e,dist)
print(min(a))
for _ in range(k):
    p = int(input())
    a = [x+i*p for i,x in enumerate(a)]
    print(min(a))

여기까지 했을 때 Pypy3 기준으로도 시관초과가 발생했다. 이는 더 적은 정점만으로 이미 더 적은 비용의 경로가 구성됐음에도 불필요하게 dist의 뒷 부분을 업데이트하고, 정점을 큐에 삽입하는 행위가 반복되기 때문인 것으로 보인다. (더 많은 수의 경로를 사용하는 것이 의미가 생기기 위해서는 비용이라도 더 적어야하는데, 그렇지 않은 경우는 이후 사용되지 않을 것이다.) 이에 효율성을 위해 탈출문 하나를 아래처럼 추가해줘야 한다.

import heapq
import sys
input=sys.stdin.readline
n,m,k=map(int,input().split())
s,e=map(int,input().split())
g=[[] for _ in range(n+1)]
for _ in range(m):
    a,b,w=map(int,input().split())
    g[a].append((b,w))
    g[b].append((a,w))
INF=int(1e9)
dist=[[INF]*(n+1) for _ in range(n+1)]
def dijk(s,e,dist):
    q = []
    heapq.heappush(q,(0,s,0))
    dist[s] = [0]*(n+1)
    while q:
        d, now, cnt = heapq.heappop(q)
        f = 0
        for i in range(1,cnt):
            if dist[now][i] < d:
                f = 1
                break
        if dist[now][cnt] < d or now==e or f:
            continue
        for x,w in g[now]:
            cost = d + w
            if cost < dist[x][cnt+1]:
                dist[x][cnt+1] = cost
                heapq.heappush(q,(cost,x,cnt+1))
    return dist[e]

a = dijk(s,e,dist)
print(min(a))
for _ in range(k):
    p = int(input())
    a = [x+i*p for i,x in enumerate(a)]
    print(min(a))

같은 정점 now를 기준으로 cnt보다 적은 개수의 경로를 사용해 더 적은 비용이 저장된 케이스가 있다면, 그 정점은 큐에서 삭제하고 넘어간다.