[python] 16236. 아기 상어

문제 출처

16236. 아기 상어

풀이

아기 상어가 물고기를 먹으며 이동하는 최단거리를 구하는 문제이다. 아기 상어는 다음과 같은 조건을 만족하는 경우에만 물고기를 먹을 수 있다.

 

1. 자신의 크기보다 작거나 같은 물고기가 있는 칸만 지나갈 수 있다.
2. 자신의 크기보다 작은 물고기만 먹을 수 있다.
3. 먹을 수 있는 물고기가 1마리보다 많다면, 거리가 가장 가까운 물고기를 먹으러 간다.
   (거리는 아기 상어가 있는 칸에서 물고기가 있는 칸으로 이동할 때, 지나야하는 칸의 개수의 최솟값이다.)
4. 거리가 가까운 물고기가 많다면, 가장 위에 있는 물고기, 그러한 물고기가 여러마리라면, 가장 왼쪽에 있는 물고기를 먹는다.

위 조건을 모두 만족하는 최단거리는 우선순위를 고려한 정렬 알고리즘을 통해 구할 수 있다. 하지만 python의 built-in library중 최소힙으로 구현된 우선순위 큐 heapq 모듈을 사용하면 비교적 간단하게 구할 수 있다.

 

 

from heapq import heappop, heappush


N = int(input())
arr = [list(map(int, input().split())) for _ in range(N)]
q = []

def init():
    for y in range(N):
        for x in range(N):
            if arr[y][x] == 9:
                heappush(q, (0, y, x))
                arr[y][x] = 0
                return

def bfs():
    global q
    body, eat, distance = 2, 0, 0
    path = [[0]*N for _ in range(N)]
    while q:
        d, y, x = heappop(q)
        if 0 < arr[y][x] < body:
            eat += 1
            arr[y][x] = 0
            if eat == body:
                body += 1
                eat = 0
            distance += d
            d = 0
            q = []
            path = [[0]*N for _ in range(N)]

        for dx, dy in (0, -1), (-1, 0), (1, 0), (0, 1):
            nd, ny, nx = d+1, y+dy, x+dx
            if nx < 0 or nx >= N or ny < 0 or ny >= N:
                continue
            if 0 < arr[ny][nx] > body or path[ny][nx]:
                continue
            path[ny][nx] = 1
            heappush(q, (nd, ny, nx))
    print(distance)

init()
bfs()