[Algorithm] 2307번 - 도로검문

문제

풀이 방법

파이썬이라는 언어의 특성 때문인지 시간초과가 정말 많이 나왔다.
그래서 시간을 단축하기 위한 소소한 팁을 적으려 한다.

1. 입력을 받을 때에는 2차원 배열(인접 행렬)이 아닌 인접 리스트로 저장한다. 다익스트라 알고리즘에서는 주로 연결된 노드만 탐색하기 때문에 인접 리스트가 더 효율적이다.

2. 다익스트라 알고리즘을 통해 경로의 최소비용을 찾을 때에는 우선순위 큐를 사용한다.

3. 도로검문을 실행하지 않을 때에는 최단 경로를 추적하기 위해서 path 배열에 노드를 저장하는 과정을 수행했다. 하지만 이후 도로검문을 수행할 때에는 해당 과정이 불필요하므로 2개의 메서드 dijkstra, modified_dijkstra로 분리하여 구현했다.

4. 도로검문을 수행할 도로는 최단 경로를 만드는 경로만 고려해도 충분하다. 최단 경로에 포함되지 않는 도로를 검문하면 지연시간이 아예 발생하지 않기 때문이다. 이를 위해 path에 저장된 노드를 역추적하며 최단 경로를 구성하는 모든 도로를 찾아냈고, 이를 check 집합에 저장했다.
   이 때, 주의할 점은 최단 경로를 만드는 방법이 여러 개가 존재할 때이다. 모든 경로를 고려하지 않는다면 6% 대에서 '틀렸습니다' 를 만날 수 있다.

tmp = [v]
check = set()
while tmp:
    now = tmp.pop()
    for before in path[now]:
        check.add((now, before))
        tmp.append(before)

5. 도로검문을 수행할 때, edges 인접 리스트에서 저장한 도로의 가중치를 INF로 바꿔서 다익스트라 알고리즘을 수행할 수도 있겠지만 이렇게 그래프를 조작하는 것은 자칫 시간초과를 유발할 수 있다. 따라서 변형된 다익스트라 알고리즘 함수modified_dijkstra에 인자로 검문하는 도로 정보를 넘겨주고, 함수 안에서는 해당 도로인 경우를 if문으로 확인하고 지나치지 않도록 continue 하는 로직을 구현했다.

max_cost = 0
for (now, before) in check:
    blocked_road = {(now, before)}
    delay = modified_dijkstra(blocked_road)
    if max_cost == float('inf'):
        break
    max_cost = max(max_cost, delay)

정답 코드

# 다익스트라 - 2307번 - 도로검문
from heapq import *
import sys
input = sys.stdin.readline


# 처음 최단경로를 찾을 때에만 사용 (경로 탐색을 위해서 배열을 하나 더 사용함)
def dijkstra():        
    result = [float('inf')] * (v + 1)
    way = [[] for _ in range(v + 1)]

    result[1] = 0
    q = []
    heappush(q, (0, 1))  # 누적 비용, 지금 노드
    # 1. 검문이 없을 때의 최소 시간 구하기
    while q:
        cost, now = heappop(q)
        if cost > result[now]:
            continue
        for next, weight in edges[now]:
            new_cost = cost + weight
            if result[next] == new_cost:
                way[next].append(now)
            if result[next] > new_cost:
                result[next] = new_cost
                way[next] = [now]
                heappush(q, (new_cost, next))
    return result, way


# 도로검문 시 사용 (경로 탐색할 필요가 없는 경우)
def modified_dijkstra(blocked_road):        
    result = [float('inf')] * (v + 1)
    result[1] = 0
    q = []
    heappush(q, (0, 1))  # 누적 비용, 지금 노드

    while q:
        cost, now = heappop(q)
        if cost > result[now]:
            continue
        for next, weight in edges[now]:
            if (now, next) in blocked_road or (next, now) in blocked_road:      # 검문을 시행하는 도로는 건너뜀
                continue
            new_cost = cost + weight
            if result[next] > new_cost:
                result[next] = new_cost
                heappush(q, (new_cost, next))
    return result[v]


v, e = map(int, input().split())
edges = [[] for _ in range(v+1)]  # 인접 리스트로 그래프 가중치를 저장
for i in range(e):
    a, b, w = map(int, input().split())
    edges[a].append((b, w))
    edges[b].append((a, w))
  
# 1. 도로검문이 없을 때 최단경로를 구한다.  (min_cost)
result, path = dijkstra()
min_cost = result[v]

# 2. 최단경로가 되는 모든 경로를 찾아서 check 집합에 저장한다. 
# (path[end] 부터 path[start]까지 역추적하며 경로를 찾아낸다)
tmp = [v]
check = set()
while tmp:
    now = tmp.pop()
    for before in path[now]:
        check.add((now, before))
        tmp.append(before)
# print(check)

# 3. check 배열에 저장된 모든 도로를 하나씩 막으며 지연시간의 최댓값을 구한다.
max_cost = 0
for (now, before) in check:
    blocked_road = {(now, before)}
    delay = modified_dijkstra(blocked_road)

    if max_cost == float('inf'):
        break
    max_cost = max(max_cost, delay)

if max_cost == float('inf'):
    print(-1)
else:
    print(max_cost - min_cost)