[백준 4485][파이썬] 녹색 옷 입은 애가 젤다지? (다익스트라)

https://www.acmicpc.net/problem/4485

4485번: 녹색 옷 입은 애가 젤다지?

 

https://growth-coder.tistory.com/200

위 포스팅에서 다익스트라 기본 문제를 풀어보았다.

 

알고리즘 자체는 거의 동일하나 보완할 필요성을 느껴서 기록을 해보려고 한다.

 

우선 이 문제는 그래프로 표현하기보다는 2차원 배열 그대로 보는 편이 좋다.

 

연결된 노드들은 동서남북으로 정해져 있기 때문에 동서남북으로 돌면서 체크하는 편이 좋다.

 

우선 이 문제는 위 포스팅에 사용한 알고리즘을 사용하면 시간 초과가 발생한다.

 

우선 시간 초과가 발생한 코드는 다음과 같다.

import sys
import heapq
input=sys.stdin.readline
MAX=1e9
dx=[0, 0, 1, -1] #동서남북 
dy=[1, -1, 0, 0] #동서남북 
def can_go(a, b, n): #n*n에서 [a][b] 갈 수 있는지
    return 0<=a<n and 0<=b<n
cnt=0
while True:
    n=int(input())
    cnt+=1
    if n==0:
        break
    arr=[
        list(map(int, input().split())) 
        for  _ in range(n)
    ]
    arr2=[
        [MAX]*(n)
        for _ in range(n)
    ]
    queue = [] #최소 힙
    #참고로 마지막에 arr[0][0] 더하는 거 잊지 말기
    heapq.heappush(queue, (0, (0, 0))) #금액과 위치
    while queue:
        cost, coordinate = heapq.heappop(queue) #금액과 위치 빼옴
        for i in range(4): #동서남북 조사
            nx=coordinate[0]+dx[i]
            ny=coordinate[1]+dy[i]
            if can_go(nx, ny, n) and cost+arr[nx][ny]<=arr2[nx][ny]: #좌표를 벗어나지 않는다면
                heapq.heappush(queue, (cost+arr[nx][ny], (nx, ny)))
                arr2[nx][ny]=cost+arr[nx][ny]
    print(f"Problem {cnt}: {arr2[n-1][n-1]+arr[0][0]}")

시간 초과가 발생한 원인은 다음 코드에 있었다.

            if can_go(nx, ny, n) and cost+arr[nx][ny]<=arr2[nx][ny]: #좌표를 벗어나지 않는다면
                heapq.heappush(queue, (cost+arr[nx][ny], (nx, ny)))
                arr2[nx][ny]=cost+arr[nx][ny]

이전 포스팅에서 cost+arr[nx][ny]<=arr2[nx][ny]와 같은 조건문을 넣은 이유를 다음과 같이 설명했었다.

 

아래 상황은 노드 A, B, C는 최단시간이 확정된 상태이다.

이 상황에서 A-C-B라는 후보 경로가 나왔는데 이는 필요가 없다. 이미 A-B라는 최단 경로가 확정된 상태이기 때문이다.

 

그래서 이를 그냥 최소 힙에 넣어두고 나중에 나오게 되면 cost+arr[nx][ny]<=arr2[nx][ny] 조건문에서 걸러지게 된다.

 

그런데 이렇게 구현하면 일단 이 쓸모없는 후보 경로가 힙에 들어가고 나오는 과정이 추가되는데 최소 힙에 값을 넣고 빼는 시간 복잡도는 O(n*log n )이다.

 

그리고 이 젤다 문제는 한 노드마다 동서남북으로 링크가 연결되어있는데 (모서리 제외) 노드마다 4번씩 (모서리 제외) 굳이 없어도 되는 최소 힙에 넣고 빼는 과정이 추가된다. 

 

그래서 시간 초과가 발생한 것이다.

 

해결 방안은 간단하다.

 

각 노드마다 방문 여부를 알려주는 2차원 배열을 생성한다.

 

그리고 최소 힙에서 값을 꺼내온다는 것은 꺼내온 노드의 최단 경로가 확정되었다는 뜻이다. (자세한 설명은 다익스트라 기본 문제 링크 참고)

 

이렇게 꺼내온 노드에는 방문을 했다는 뜻인 True를 저장하고 만약 해당 노드를 방문했다면 (해당 노드의 최단 경로가 확정된 상태라면) 최소 힙에 아예 넣지를 않는 것이다.

 

이렇게 쓸모없이 최소 힙에 들어갔다 나오는 과정을 제거하여 제한 시간 내에 통과할 수 있었다.

 

<최종 코드>

import sys
import heapq
input=sys.stdin.readline
MAX=1e9
dx=[0, 0, 1, -1] #동서남북 
dy=[1, -1, 0, 0] #동서남북 
def can_go(a, b, n): #n*n에서 [a][b] 갈 수 있는지
    return 0<=a<n and 0<=b<n
cnt=0
while True:
    n=int(input())
    cnt+=1
    if n==0:
        break
    arr=[
        list(map(int, input().split())) 
        for  _ in range(n)
    ]
    arr2=[
        [MAX]*(n)
        for _ in range(n)
    ]
    visited=[
        [False]*(n)
        for _ in range(n)
    ]
    queue = [] #최소 힙
    #참고로 마지막에 arr[0][0] 더하는 거 잊지 말기
    heapq.heappush(queue, (0, (0, 0))) #금액과 위치
    while queue:
        cost, coordinate = heapq.heappop(queue) #금액과 위치 빼옴
        for i in range(4): #동서남북 조사
            nx=coordinate[0]+dx[i]
            ny=coordinate[1]+dy[i]
            if can_go(nx, ny, n) and not visited[nx][ny]: #좌표를 벗어나지 않는다면
                heapq.heappush(queue, (cost+arr[nx][ny], (nx, ny)))
                arr2[nx][ny]=cost+arr[nx][ny]
                visited[nx][ny]=True
    print(f"Problem {cnt}: {arr2[n-1][n-1]+arr[0][0]}")