[프로그래머스] 미로 탈출

미로 탈출

풀이

1. 간선이 입력될 때, 같은 것이 여러번 입력될 수 있다.
   즉 최단 거리를 구하기 위해선 거리가 가장 짧은 것을 골라내야 한다.
   그래프는 배열로 이루어져 있으며 노드의 최대 크기 n만큼의 원소를 가진다.
   i번째 노드의 정보를 담고 있는 graph[i]는
   목적지를 키로, 거리를 값으로 하는 딕셔너리다.
   엣지를 입력받으면서 이 값들을 항상 최소값으로 유지한다.

2. 갈 때마다 노드에 딸린 간선의 방향이 바뀐다.
   즉 현재 함정이 어떤 것이 켜져있는지를 알아야 한다.
   그래야 내가 A에서 B를 갈 때 그 간선의 방향을 알 수 있다.

3. 함정이 켜진 노드들을 추적할 수 있다면 다익스트라 알고리즘을 적용할 수 있다.
   다익스트라를 돌리면서 현재 정점 now_node와 내가 갈 곳 adj의 함정 여부를 알아야 한다.
   둘 다 함정이 켜져있거나 꺼져있다면 (= 둘의 상태가 같다면) 정방향으로 가야 한다.
   둘의 상태가 다르면 역방향으로 가야 한다.

4. 정방향 엣지를 쉽게 알아내기 위해서, 그래프를 입력받을 때 역방향은 음수로 표시한다.
   예를 들어 5번 노드에서 3번 노드로 가는 거리 3짜리 엣지가 있다고 하면
   3번 노드는 -5번 노드로 가는 거리가 똑같은 엣지가 있다고 볼 수 있다.

5. 현재 함정 활성화 상태를 추적하면서 거리를 최소화 해나가야 한다.
   즉 dist 배열에서 start부터 노드 v까지의 거리를 구할 때,
   그냥 dist[v]가 아니라 dist[v][start]로 구해야 한다.

6. 일반적인 다익스트라에서 몇 가지 다른 부분을 가진다.

6-1.
먼저 start에서 도달할 수 있는 노드까지들의 거리를 저장하는 dist 배열은
당시의 함정 상황을 같이 기록하기 위해서 2차원 배열로 바뀐다.
처음에 해당 노드의 on/off를 기록할 수 있는 공간을 둬서 2차원을 2의 크기로 잡았는데,
이렇게 하면 나중에 end로 도달하지 못하는 예가 생긴다.
2 ** 10로 해서 모든 경우의 수를 같이 저장할 수 있도록 해야 한다.
상태 추적이 추가되었다.

6-2.
우선순위 큐에서 뽑아낸 현재 노드에 달린 간선들을 살피며 앞으로 나아갈 때,
3번과 4번에 의해서 지금 가려고 하는 노드가 음수인지 양수인지 알아야 한다.
노드가 음수인데 (역방향) 현재 엣지도 역방향이라면 갈 수 있고 정방향이면 갈 수 없다.
즉, 기존 다익스트라에서 갈 수 있는지 없는지 판단하는 작업이 하나 추가된 것이다.

코드

import heapq
import math
from collections import defaultdict


def solution(n, start, end, roads, traps):
    graph = make_graph(n, roads)
    traps = {traps[i]: i for i in range(len(traps))}
    dist = dijkstra(start, graph, traps)
    return min(dist[end])


def make_graph(n, roads):
    graph = [defaultdict(lambda: 3001) for _ in range(n + 1)]
    for P, Q, S in roads:
        graph[P][Q] = min(graph[P][Q], S)
        graph[Q][-P] = min(graph[Q][-P], S)
    return graph


def dijkstra(start, graph, traps):
    dist = [[math.inf for _ in range(2 ** 10)] for _ in range(len(graph))]
    dist[start][0] = 0
    pq = [(0, start, 0)]

    while len(pq) > 0:
        now_dist, now_node, now_status = heapq.heappop(pq)

        for adj, adj_dist in graph[now_node].items():
            direction = True if adj > 0 else False
            adj = abs(adj)

            adj_trapped = 0 if adj not in traps.keys() else ((now_status & (1 << traps[adj])) != 0)
            now_trapped = 0 if now_node not in traps.keys() else ((now_status & (1 << traps[now_node])) != 0)
            edge_reversed = adj_trapped != now_trapped

            if direction != edge_reversed:
                next_status = now_status

                if adj in traps:
                    next_status = now_status ^ (1 << traps[adj])

                new_dist = now_dist + adj_dist
                if dist[adj][next_status] > new_dist:
                    dist[adj][next_status] = new_dist
                    heapq.heappush(pq, (new_dist, adj, next_status))

    return dist