[파이썬] 프로그래머스 : 미로 탈출 (Lv.4)

[파이썬] 프로그래머스 : 미로 탈출 (Lv.4)

프로그래머스

---

풀이

0.방향성 생각

• node가 1000개, edge가 3000개로 입력이 작은 경우.
• trap수가 하필 10개. -> bitmask?
• state 2^10 * visit node 1000 = 10만
• TSP처럼 dfs+dp+bitmask로 풀든가, 다익+bitmask로 풀기 -> 재귀는 실수할거같아서 다익 + bitmask
• 현재 노드 x와 다음 노드 nx가 있을 때, trap 유무로 4가지 케이스로 크게 분류한다.

1. 연결 리스트, 간선 cost 저장.

import heapq as hq
from collections import defaultdict as dd

def solution(n, start, end, roads, traps):
    
    inf = 1e9
    trap = {val:idx for idx,val in enumerate(traps)}
    costs = {i: {j:inf for j in range(1,n+1)} for i in range(1,n+1)}
    graph = [[] for _ in range(n+1)]
    for p,q,cost in roads:
        costs[p][q] = min(costs[p][q],cost) # 초기 상태 방향중에 가장 작은값으로, 양방향 간선이 아님
        graph[p].append(q)
        graph[q].append(p)

• costs에는 p -> q방향중 가장 작은 값을 가지는 간선을 선택한다.
• 양방향 간선이 아닌데도 graph에 양방향으로 추가해주는 이유는 trap 상태에 따라 연결유무가 바뀔수도 있어서, 일단 접근을 한 후에, costs에서 정의된 cost를 보고 간선 정보를 받아올 수 있음.

2. 탐색

    visit = [[inf]*(n+1) for _ in range(1<<len(traps))]
    visit[0][start] = 0
    
    heap = [(0,start,0)] # cost,시작,state
    while heap:
        c,x,s = hq.heappop(heap)
        
        if x == end:
            return c

		for nx in graph[x]:

• visit 배열 같은 경우는 1<<트랩개수로 10개면 2^10 = 1024이다.
• visit[state][node] 순으로 저장
• 힙에 cost,node,state 순으로 저장하고 탐색 시작
• x가 end에 도착하면 탈출하기
• 연결리스트 graph에서 연결정보 x와 연결된 nx의 정보를 받아온다.

3. x(함정x) -> nx(함정x)

            if x not in trap and nx not in trap: # 둘다 함정 아니면
                if c+costs[x][nx] < visit[s][nx]:
                    nc = c+costs[x][nx]
                    hq.heappush(heap,(nc,nx,s))
                    visit[s][nx] = nc

• x,nx가 둘다 함정이 아닌 경우, 간선의 정보가 바뀔 염려가 없다.
• 그냥 기존 다익스트라처럼 풀면 된다.

4. x(함정o) -> nx(함정x)

            elif x in trap and nx not in trap: # 함정o -> 함정x
                if costs[nx][x] != inf and s&1<<trap[x] and c+costs[nx][x] < visit[s][nx]:
                    nc = c+costs[nx][x] # 함정 켜저서 간선 뒤집힌경우
                    hq.heappush(heap,(nc,nx,s))
                    visit[s][nx] = nc
                elif not s&1<<trap[x] and c+costs[x][nx] < visit[s][nx]:
                    nc = c+costs[x][nx] # 함정 안켜저서 그냥 가는경우
                    hq.heappush(heap,(nc,nx,s))
                    visit[s][nx] = nc

• 현재 함정의 on/off 유무로 간선 정보가 뒤집혀있는지 확인한다.
• 함정이 켜져있으면 뒤집힌 간선을 타고 이동한다. costs[nx][x]가 연결되었는지 확인한다.
• 비트마스킹을 통해 현재 node x의 함정이 켜저있는지 확인하고, 켜진 경우 갱신하는경우에 업데이트.
• 함정이 꺼져있는 경우, graph[x]에서 이미 cost가 inf가 아님을 확인했으므로, 바로 갱신하는 경우에 업데이트.

5. x(함정x) -> nx(함정o)

            elif x not in trap and nx in trap: # 함정x -> 함정o
                ns = s^1<<trap[nx]
                if costs[nx][x] != inf and s>>trap[nx]&1 and c+costs[nx][x] < visit[ns][nx]: # 함정 켜저서 간선 뒤집힌경우
                    nc = c+costs[nx][x]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc
                    
                elif not s>>trap[nx]&1 and c+costs[x][nx] < visit[ns][nx]: # 함정 꺼져서 그냥 가는경우
                    nc = c+costs[x][nx]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc

• nx로 이동 시, 함정을 밟으면 상태가 바뀐다. nx위치의 함정을 토글해주면서 다음 상태를 ns로 정의한다.
• 마찬가지로 함정이 켜진 유무에 따라서 간선 정보를 확인하고 갱신이 가능하면 힙에 추가한다.

6. x(함정o) -> nx(함정o)

            else:
                ns = s^1<<trap[nx]
                if (s>>trap[x]&1)^(s>>trap[nx]&1) and costs[nx][x] != inf and c+costs[nx][x] < visit[ns][nx]: # 홀수번이라 뒤집힘
                    nc = c+costs[nx][x]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc
                elif not (s>>trap[x]&1)^(s>>trap[nx]&1) and c+costs[x][nx] < visit[ns][nx]:
                    nc = c+costs[x][nx]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc

• x와 nx가 모두 함정인 경우 두 함정이 둘다 켜지거나, 둘다 꺼지면 원상태의 간선을 타고, 그렇지 않은 경우에는 뒤집힌 간선을 타게된다. XOR를 통해서 같은경우에는 0, 다른 경우에는 1로 구한다.
• 간선이 뒤집히지 않은 경우와 뒤집힌 경우 같은 방식으로 풀이를 진행한다.

 

---

전체코드

import heapq as hq
from collections import defaultdict as dd

def solution(n, start, end, roads, traps):
    
    inf = 1e9
    trap = {val:idx for idx,val in enumerate(traps)}
    costs = {i: {j:inf for j in range(1,n+1)} for i in range(1,n+1)}
    graph = [[] for _ in range(n+1)]
    for p,q,cost in roads:
        costs[p][q] = min(costs[p][q],cost) # 초기 상태 방향중에 가장 작은값으로, 양방향 간선이 아님
        graph[p].append(q)
        graph[q].append(p)

    visit = [[inf]*(n+1) for _ in range(1<<len(traps))]
    visit[0][start] = 0

    heap = [(0,start,0)] # cost,시작,state
    while heap:
        c,x,s = hq.heappop(heap)

        if x == end:
            return c

        for nx in graph[x]:

            if x not in trap and nx not in trap: # 둘다 함정 아니면
                if c+costs[x][nx] < visit[s][nx]:
                    nc = c+costs[x][nx]
                    hq.heappush(heap,(nc,nx,s))
                    visit[s][nx] = nc
                    
            elif x in trap and nx not in trap: # 함정o -> 함정x
                if costs[nx][x] != inf and s&1<<trap[x] and c+costs[nx][x] < visit[s][nx]:
                    nc = c+costs[nx][x] # 함정 켜저서 간선 뒤집힌경우
                    hq.heappush(heap,(nc,nx,s))
                    visit[s][nx] = nc
                elif not s&1<<trap[x] and c+costs[x][nx] < visit[s][nx]:
                    nc = c+costs[x][nx] # 함정 안켜저서 그냥 가는경우
                    hq.heappush(heap,(nc,nx,s))
                    visit[s][nx] = nc   
                    
            elif x not in trap and nx in trap: # 함정x -> 함정o
                ns = s^1<<trap[nx]
                if costs[nx][x] != inf and s>>trap[nx]&1 and c+costs[nx][x] < visit[ns][nx]: # 함정 켜저서 간선 뒤집힌경우
                    nc = c+costs[nx][x]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc
                    
                elif not s>>trap[nx]&1 and c+costs[x][nx] < visit[ns][nx]: # 함정 꺼져서 그냥 가는경우
                    nc = c+costs[x][nx]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc
                    
            else:
                ns = s^1<<trap[nx]
                if (s>>trap[x]&1)^(s>>trap[nx]&1) and costs[nx][x] != inf and c+costs[nx][x] < visit[ns][nx]: # 홀수번이라 뒤집힘
                    nc = c+costs[nx][x]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc
                elif not (s>>trap[x]&1)^(s>>trap[nx]&1) and c+costs[x][nx] < visit[ns][nx]:
                    nc = c+costs[x][nx]
                    hq.heappush(heap,(nc,nx,ns))
                    visit[ns][nx] = nc

코멘트

코드가 좀 길긴 한데, 거의 복붙이라서 금방 풀었다.

복붙하다가 고치는 부분 빼먹어서 약간 삽질은 했는데 만족.

다음에 풀때는 dfs + dp로 풀어보기