백준 / 아기 상어 / 16236번 / Python

*문제 출처는 백준에 있습니다.

문제 제목: 아기 상어 / 16236번 (골드 3단계)

문제 사이트: https://www.acmicpc.net/problem/16236

 

문제 설명

---

나의 풀이

from collections import deque

def solution(ftank, n):
    # BFS 알고리즘을 사용할 큐
    q = deque()

    # 걸리는 시간
    max_time = 0

    # 움직이기
    move = [(0, 1), (0, -1), (1, 0), (-1, 0)]

    for i in range(n):
        for j in range(n):
            if ftank[i][j] == 9:
                # 아기 상어 위치 i, j / size = 2 / 먹이 먹은 횟수 = 0 / time = 0
                q.append([i, j, 2, 0, 0])
                ftank[i][j] = 0

    while q:
        # 상어 사이즈를 사용할 변수 및 먹은 개수
        # 초기 사이즈는 2이며 사이즈만큼 먹어야 커진다.
        y, x, shark_size, shark_eat, time = q.popleft()

        max_time = max(max_time, time)

        for iy, ix in move:
            ny = iy + y
            nx = ix + x
            if 0 <= ny < n and 0 <= nx < n:
                if ftank[ny][nx] == 0 or ftank[ny][nx] == shark_size:
                    # 빈칸이거나 크기가 같은 물고기라면 이동 가능
                    q.append([ny, nx, shark_size, shark_eat, time + 1])
                elif 0 < ftank[ny][nx] < shark_size:
                    # 상어보다 작은 물고기를 먹을 수 있는 경우
                    ftank[ny][nx] = 0
                    if shark_eat + 1 == shark_size:
                        # 상어가 자신의 크기만큼 먹으면 크기 증가
                        q.append([ny, nx, shark_size + 1, 0, time + 1])
                    else:
                        q.append([ny, nx, shark_size, shark_eat + 1, time + 1])

        # 무한 루프를 빠져나가기 위한 변수
        end = False

        for i in range(n):
            for j in range(n):
                if 0 < ftank[i][j] < shark_size:
                    end = True
                    break
            if end:
                break

        if end:
            continue
        else:
            max_time += 1
            break

    return max_time

N = int(input())

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

print(solution(shark, N))

정답은 나오지만 시간초과가 발생하는 코드이다. BFS 탐색에서 비효율적으로 모든 경우를 계속 탐색하고 있기 때문이다. 특히, 먹을 수 있는 물고기가 없더라도 계속해서 탐색을 진행하는 로직이 있기 때문에 시간이 낭비되고 있었다.

그래서 BFS에서 우선순위 큐를 사용하고, 한 번의 BFS 탐색으로 가장 가까운 먹이를 탐색하는 방법을 사용했다.

import heapq
from collections import deque

# 상하좌우 이동을 위한 방향 설정
move = [(-1, 0), (1, 0), (0, -1), (0, 1)]

# BFS로 먹을 수 있는 가장 가까운 물고기 찾기
def bfs(shark, n, shark_size):
    q = deque([shark])  # 아기 상어의 위치를 시작점으로 BFS
    visited = [[False] * n for _ in range(n)]
    visited[shark[0]][shark[1]] = True
    dist = [[-1] * n for _ in range(n)]  # 거리 저장
    dist[shark[0]][shark[1]] = 0

    fish_list = []  # 먹을 수 있는 물고기 후보들

    while q:
        y, x = q.popleft()

        for iy, ix in move:
            ny, nx = y + iy, x + ix

            if 0 <= ny < n and 0 <= nx < n and not visited[ny][nx]:
                if ftank[ny][nx] <= shark_size:  # 상어가 지나갈 수 있는 칸
                    visited[ny][nx] = True
                    dist[ny][nx] = dist[y][x] + 1
                    q.append((ny, nx))

                    # 먹을 수 있는 물고기를 발견
                    if 0 < ftank[ny][nx] < shark_size:
                        heapq.heappush(fish_list, (dist[ny][nx], ny, nx))

    if fish_list:
        # 가장 가까운 물고기를 반환 (거리가 가까운, 가장 위, 가장 왼쪽)
        return heapq.heappop(fish_list)
    return None

def solution(ftank, n):
    # 아기 상어 초기 상태
    for i in range(n):
        for j in range(n):
            if ftank[i][j] == 9:
                shark = (i, j)  # 상어의 위치
                ftank[i][j] = 0  # 상어가 있는 칸을 빈 칸으로

    shark_size = 2  # 초기 크기
    shark_eat = 0  # 먹은 물고기 수
    time = 0  # 걸린 시간

    while True:
        # BFS로 가장 가까운 먹을 수 있는 물고기 찾기
        result = bfs(shark, n, shark_size)

        if result is None:  # 더 이상 먹을 물고기가 없으면 종료
            break

        # 물고기를 먹고, 상어의 위치와 시간 갱신
        dist, ny, nx = result
        time += dist  # 이동한 거리가 시간
        shark = (ny, nx)
        ftank[ny][nx] = 0  # 물고기를 먹은 자리는 빈 칸으로

        # 상어가 물고기를 먹으면 크기 증가 여부 확인
        shark_eat += 1
        if shark_eat == shark_size:
            shark_size += 1
            shark_eat = 0

    return time

# 입력 받기
N = int(input())
ftank = [list(map(int, input().split())) for _ in range(N)]

# 결과 출력
print(solution(ftank, N))

우선순위 큐를 사용: 먹을 수 있는 물고기를 거리 순서, 위쪽 우선, 왼쪽 우선으로 자동 정렬하여 가장 적합한 물고기를 먼저 찾는다. BFS 최적화를 통해서 한 번에 가장 가까운 먹을 수 있는 물고기를 탐색한다.

---

※ 알아야 할 것

https://newkimjiwon.tistory.com/77

[Algorithm] 너비 우선 탐색(BFS, Breadth-First Search)

https://newkimjiwon.tistory.com/179

[자료구조] 힙(Heap)