[백준] 16517 | Tima goes to Xentopia [파이썬/python]

문제

link: https://www.acmicpc.net/problem/16517

시작점 S에서 T까지 가는 가장 짧은 시간을 출력하는 문제이다. s에서 t로 이동할 때는 정확히 k1개의 빨간색 선로와 k2개의 파란색 선로를 사용해야 한다. 불가능하다면 -1을 출력한다.

 

접근

우선순위 큐를 활용했고, dist[v][r][b]에 r개의 빨간색 선로와 b개의 파란색 선로를 사용해서 노드 v에 도달했을 때 걸리는 시간의 최솟값을 기록하도록 했다. 

 

코드

import sys
input = sys.stdin.readline
import heapq

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

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

s, t = map(int, input().split())

dist = [[[float('inf')]*(k2+1) for _ in range(k1+1)] for _ in range(n+1)]
h = []
dist[s][0][0] = 0
heapq.heappush(h, (0, s, 0, 0))

while h:
    x, v, r, b = heapq.heappop(h)
    if dist[v][r][b] < x:
        continue

    for nv, nx, nc in graph[v]:
        nr, nb = r, b
        if nc == 1:
            nr += 1
        elif nc == 2:
            nb += 1

        if nr <= k1 and nb <= k2 and dist[nv][nr][nb] > x+nx:
            dist[nv][nr][nb] = x+nx
            heapq.heappush(h, (x+nx, nv, nr, nb))

print(dist[t][k1][k2] if dist[t][k1][k2] != float('inf') else -1)