화성 탐사 코드 및 해설 (파이썬)

출처: 이것이 취업을 위한 코딩 테스트다 with 파이썬, Q 39, p 388

 

다익스트라 알고리즘을 활용해 풀이했습니다.
우선 입력을 받아 하나의 테스트 케이스는 N*N 리스트로 정리하고, 각 테스트 케이스를 spaces라는 리스트에 저장했습니다.
화성 탐사 기계는 상하좌우 인접한 곳으로 1칸씩 이동할 수 있다고 했으므로, 그래프의 인접한 노드 대신 상하좌우의 좌표를 탐색하며 풀이했습니다. 이때 N*N 리스트 밖으로 나가 IndexError가 발생하는 것을 방지하기 위해 인덱스가 0~N-1 사이인지 확인해주었습니다. 마지막으로 가장 위 칸인 (0,0)을 지나는데도 에너지가 소모되므로, 다익스트라 알고리즘이 끝난 후에 (0,0)을 지나는데 소모되는 에너지를 더해주었습니다.

 

import fileinput
import heapq

INF = int(1e9)

# 입력 처리하기
nextN = 1
spaces = []
for i, line in enumerate(fileinput.input()):
    if i == 0:
        continue
    elif i == nextN:
        N = int(line)
        spaces.append([])
        nextN += N + 1
    else:
        spaces[-1].append(list(map(int, line.split())))

# 다익스트라 알고리즘
def explore(space):
    N = len(space)
    # N * N
    distance = [[INF]*N for _ in range(N)]
    q = []
    # (0,0)에서 시작
    heapq.heappush(q, (0, (0,0)))
    distance[0][0] = 0
    while q:
        dist, now = heapq.heappop(q)
        x,y = now
        if distance[x][y] < dist:
            continue
        # 상, 하, 좌, 우
        adjs = [(x-1, y), (x+1, y), (x, y-1), (x, y+1)]
        for adj in adjs:
            a, b = adj
            if a in range(0,N) and b in range(0,N):
                cost = dist + space[a][b]
                if cost < distance[a][b]:
                    distance[a][b] = cost
                    heapq.heappush(q, (cost, (a,b)))

    # (N-1, N-1)까지 가는데 소모되는 에너지 + (0,0)을 지나는데 소모되는 에너지
    return distance[N-1][N-1] + space[0][0]

for space in spaces:
    print(explore(space))