백준: 16236 아기 상어

문제

16236번: 아기 상어

 

문제 풀이

 아기 상어가 NxN 크기의 공간에 물고기를 먹으면서 이동 할때, 최단 시간을 반환하여야 한다. 문제에서는 아기 상어가 움직이기 위한 조건들이 존재한다.

 

• 아기 상어는 자신보다 크거나 같은 크기의 물고기는 먹을 수 없다.
• 가장 가까이 있는 물고기 부터 우선적으로 먹어야 한다.
• 여러 마리를 먹을 수 있다면, 가장 왼쪽에 있는 물고기를 우선 순위로 먹어야 한다.

 

 문제를 처음에는 `BFS`와 `sort`를 통해 풀 었지만, `heapq`를 활용할 경우 더욱 깔끔하고 코드가 직관적이다. `heapq`를 사용하여 우선 순위를 `(시간, x, y)`로 하여 탐색을 진행한다. 조건에 따라 더 이상 먹을 수 있는 물고기가 없는 경우 종료하면 문제에서 원하는 답을 찾을 수 있다.

 

코드

from sys import stdin
from heapq import heappush, heappop


def visitable(x, y, size, visited):
    return 0 <= x < n and 0 <= y < n and graph[x][y] <= size and not visited[x][y]


def bfs(time, shark):
    size, eat, answer = 2, 0, 0

    heap = []
    heappush(heap, (time, shark[X], shark[Y]))
    visited = [[False] * n for _ in range(n)]

    while heap:
        time, x, y = heappop(heap)

        # 먹을 수 있는 물고기라면
        if 0 < graph[x][y] < size:
            eat += 1
            graph[x][y] = 0
            # 크기가 2라면 2마리를 먹어야 함.
            if eat == size:
                size += 1
                eat = 0

            answer += time
            time = 0

            # heap, 방문 여부 초기화
            heap = []
            visited = [[False] * n for _ in range(n)]

        for dx, dy in dirs:
            next_x, next_y = x + dx, y + dy

            if visitable(next_x, next_y, size, visited):
                visited[next_x][next_y] = True
                heappush(heap, (time + 1, next_x, next_y))

    return answer


if __name__ == '__main__':
    X, Y = 0, 1
    n = int(stdin.readline())
    graph = [list(map(int, stdin.readline().split())) for _ in range(n)]
    dirs = ((0, 1), (1, 0), (0, -1), (-1, 0))

    for i in range(n):
        for j in range(n):
            if graph[i][j] == 9:
                graph[i][j] = 0
                print(bfs(0, [i, j]))
                break