백준 16236번 : 아기 상어 in Python

16236번: 아기 상어

 

www.acmicpc.net

이 문제는 여러 마리의 물고기와 한 마리의 아기 상어가 있는 공간에서 아기 상어가 자신이 먹을 수 있는 물고기를 다 먹을 때까지 걸리는 시간을 구하는 문제이다. 이때 아기 상어는 자신의 크기보다 큰 물고기가 있는 칸은 못 지나가고, 자신의 크기보다 작은 물고기는 먹을 수 있다. 그리고 먹을 수 있는 물고기들 중 항상 거리가 가까운 물고기를 먹으러 간다. 이때 거리가 가까운 물고기가 많은 경우 가장 위에 있는 물고기, 가장 왼쪽에 있는 물고기 순으로 먹는다.

우선 매번 물고기를 먹은 뒤에 아기 상어의 위치를 기준으로 가장 가까운 물고기를 찾아야 하므로 이를 구현하는 함수를 만들어야 한다. 이때 이를 찾기 위해서는 BFS를 실행해야 하므로, 따라서 상어의 좌표를 매개변수로 하는 BFS를 구현하는 함수를 만들어야 한다. 이때 주의해야 할 점은 그래프 탐색을 할 때 먹을 수 있는 물고기나 지나갈 수 있는 물고기는 상어의 크기를 기준으로 정하기 때문에 함수의 매개변수로 상어의 크기도 주어야 한다. 이제 BFS를 함수 내에서 구현하는데, 이때 먹을 수 있는 물고기 중 가장 가까운 물고기를 골라야 하므로 일단 상어 위치를 기준으로 먹을 수 있는 물고기를 배열에 저장한 뒤, 가장 가까운 물고기를 선택한다.

이제 상어의 초기 좌표를 시작으로 BFS를 실행한 뒤 가장 가까운 물고기를 얻어 해당 좌표의 값을 0으로 바꾼 뒤, 상어가 지금까지 먹은 물고기 수를 1 증가시킨다. 이때 만약 지금까지 먹은 물고기 수가 상어의 크기와 같은 경우 상어의 크기를 하나 늘리고 지금까지 먹은 물고기 수를 0으로 초기화한다. 그리고 그 물고기까지 이동하는데 걸린 시간을 결과에 추가하고, 먹은 물고기의 좌표에서 다시 위 과정을 반복한다. 이를 먹을 수 있는 물고기가 없을 때까지 반복하면 된다.

이를 코드로 구현하면 다음과 같다.

import heapq
from collections import deque
N = int(input())
space = []
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]
for _ in range(N):
    space.append(list(map(int, input().split())))
# 상어가 먹을 수 있는 물고기가 없을 때까지 걸리는 시간을 구하는 함수
def shark_move(start_x, start_y, ate_num, size_shark):
    # 상어가 현재 먹을 수 있는 물고기를 저장하는 배열
    can_eat = []
    # BFS를 실행해 현재 상어 위치에서 먹을 수 있는 물고기를 can_eat 배열에 저장한다.
    bfs(start_x, start_y, size_shark, can_eat)
    # 먹을 수 있는 물고기가 없으면 더 이상 이동할 필요가 없으므로 0을 반환한다.
    if len(can_eat) == 0: return 0
    # can_eat를 min heap 형태로 구현해 거리가 가까운 물고기, 가장 위에 있는 물고기, 
    # 가장 왼쪽에 있는 물고기 순서대로 선택한다.
    d, x, y = heapq.heappop(can_eat)
    # 해당 좌표에 있는 물고기를 먹는다.
    space[x][y] = 0
    # 지금까지 먹은 물고기 수를 1 늘린다.
    ate_num += 1
    # 지금까지 먹은 물고기 수와 상어 크기가 같으면 상어 크기를 1 늘리고 
    # 지금까지 먹은 물고기 수를 초기화한다.
    if ate_num == size_shark:
        ate_num = 0
        size_shark += 1
    # 이제 현재 좌표에 대해 다시 이 함수를 실행해 그 결과에 지금까지 걸린 시간을 더한다.
    return d + shark_move(x, y, ate_num, size_shark)
# 상어의 위치를 기준으로 BFS를 실행하는 함수
def bfs(i, j, size_shark, can_eat):
    global N
    visited = [[False]*N for _ in range(N)]
    queue = deque()
    # 좌표와 해당 좌표까지 이동하는데 걸리는 시간을 같이 저장한다.
    queue.append([i, j, 0])
    while queue:
        x, y, d = queue.popleft()
        if visited[x][y]: continue
        visited[x][y] = True
        for i in range(4):
            next_x = x + dx[i]
            next_y = y + dy[i]
            # 상어의 크기가 물고기의 크기보다 크거나 같을 때만 지나갈 수 있다.
            if 0<=next_x<N and 0<=next_y<N and not visited[next_x][next_y] and space[next_x][next_y] <= size_shark:
                queue.append([next_x, next_y, d+1])
                # 좌표의 물고기의 크기가 상어의 크기보다 작을 때만 물고기를 먹을 수 있다.
                if 0 < space[next_x][next_y] < size_shark:
                    heapq.heappush(can_eat, [d+1, next_x, next_y])
shark = [0, 0]
# 상어의 위치를 찾고 해당 좌표의 수를 0으로 해 나중에 이 좌표를 지나갈 수 있도록 한다.
for i in range(N):
    for j in range(N):
        if space[i][j] == 9:
            shark = [i, j]
            space[i][j] = 0
print(shark_move(shark[0], shark[1], 0, 2))