[프로그래머스] 미로 탈출(Python)

https://programmers.co.kr/learn/courses/30/lessons/81304

코딩테스트 연습 - 미로 탈출

방문처리에서 굉장히 헷갈렸던 문제다. 다른 풀이를 보고 참고했다.

비트 마스크에 대해 알게 되었다.

 

import heapq
INF = 987654321

def dijkstra(n,graph,start,end,traps):
    visited = [[False for _ in range(1<<len(traps))] for _ in range(n+1)]
    q = list()
    heapq.heappush(q,(0,start,0))
    while q:
        cost, cur, state = heapq.heappop(q)
        if cur == end:
            return cost
        if visited[cur][state]:
            continue
        visited[cur][state] = True
        curTrapped = False
        trapped = {}
        
        for i in range(len(traps)):
            bit = 1 << i
            if state & bit: # 해당 함정이 밟혀있는 상태
                if cur == traps[i]: # 밟은 곳이 그 함정이라면
                    state &= ~bit # 상태 값 초기화
                else:
                    trapped[traps[i]] = True
            else:
                if cur == traps[i]:
                    state |= bit
                    trapped[traps[i]] = True
                    curTrapped = True
        for v in range(1,n+1):
            if v == cur:
                continue
            nextTrapped = True if v in trapped else False
            if nextTrapped == curTrapped: # False False 이거나 True True 이면
                if graph[cur][v] != INF:
                    heapq.heappush(q,(cost+graph[cur][v],v,state))
            else:
                if graph[v][cur] != INF:
                    heapq.heappush(q,(cost+graph[v][cur],v,state))
    return INF
        

def solution(n, start, end, roads, traps):
    graph = [[INF for _ in range(n+1)] for _ in range(n+1)]
    for i in range(1,n+1):
        graph[i][i] = 0
    for road in roads:
        p, q, s = road
        if graph[p][q] > s:
            graph[p][q] = s
    return dijkstra(n,graph,start,end,traps)