Baekjoon 16236. 아기 상어 / Python

알고리즘 설명

문제 설명

주어진 크기의 바다에서 아기 상어가 먹을수 있는 물고기를 찾고, 그 과정을 반복하여 최대한 오랜 시간동안 먹이를 먹게 해야한다.
자신보다 작은 물고기만 먹을 수 있고, 자신과 크기가 같은 물고기는 지나갈 수 있다.
먹을 수 있는 물고기가 있으면 먹으러가고, 그 수가 두 마리 이상이라면, 가장 가까운 물고기, 위쪽의 물고기, 왼쪽의 물고기 순으로 물고기를 선택한다.

코드 설명

1. 가능한 먹이를 찾는다.
2. 먹을 수 있는 물고기들을 거리, 위쪽, 왼쪽 기준으로 정렬해서 목표를 찾는다.
3. 먹이를 먹으면 상어의 크기를 증가시키고 또한 시간을 누적한다.

bfs 호출 시마다 모든 물고기를 확인하므로 시간 복잡도는 O(time * n ^ 2)이고, 공간복잡도는 O(3 * n ^ 2)이다.

from collections import deque as dq  
input = sys.stdin.readline  

n = int(input())  # 바다의 크기
sea = []  # 바다 상태

# 바다 상태 입력 
for i in range(n):
    tmp = list(map(int, input().split())) 
    for j in range(n):
        if tmp[j] == 9:  # 상어의 초기 위치
            start = (i, j)  
    sea.append(tmp)  # 바다 상태 저장

# 상하좌우 방향
direction = [(-1, 0), (0, -1), (0, 1), (1, 0)]

# BFS 알고리즘을 통해 상어가 먹을 수 있는 물고기들을 찾는 함수
def bfs(s, size):
    visited = [[0] * n for _ in range(n)]  # 방문한 칸을 기록할 리스트
    visited[s[0]][s[1]] = 1  # 상어의 시작 위치를 방문 처리
    time = [[0] * n for _ in range(n)]  # 각 칸까지 도달하는 시간 기록
    q = dq([s])  # BFS를 위한 큐 초기화
    tmp = []  # 먹을 수 있는 물고기들을 저장할 리스트

    while q:
        now_x, now_y = q.popleft()  # 현재 상어의 위치
        for i in range(4):
            nx = now_x + direction[i][0]  # 다음 위치의 x 좌표
            ny = now_y + direction[i][1]  # 다음 위치의 y 좌표
            if 0 <= nx < n and 0 <= ny < n and not visited[nx][ny]:
                if sea[nx][ny] <= size:  # 이동 가능한 칸인지 확인
                    q.append((nx, ny))  # 이동 가능한 칸을 큐에 추가
                    visited[nx][ny] = 1  # 방문 처리
                    time[nx][ny] = time[now_x][now_y] + 1  # 도달 시간 기록
                    if 0 < sea[nx][ny] < size:  # 먹을 수 있는 물고기인지 확인
                        tmp.append([nx, ny, time[nx][ny]])  # 먹을 수 있는 물고기 리스트에 추가
    # 먹을 수 있는 물고기들을 거리와 위치 기준으로 정렬
    return sorted(tmp, key=lambda x: (x[2], x[0], x[1]))

res = 0  # 총 걸린 시간을 저장할 변수
large = 2  # 상어의 초기 크기
cnt = 0  # 먹은 물고기의 수
sea[start[0]][start[1]] = 0  # 상어의 초기 위치를 빈 칸으로 변경

# BFS를 반복하여 먹이를 찾고 먹는 과정
while True:
    shark = bfs(start, large)  # 현재 크기로 BFS를 실행하여 먹이를 찾음
    if not shark:  # 더 이상 먹을 수 있는 물고기가 없다면 종료
        break
    x, y, time_ = shark.pop()  # 먹을 물고기를 선택 (가장 가까운 물고기)
    sea[x][y] = 0  # 선택한 물고기를 먹음
    res += time_  # 총 시간에 이동 시간 추가
    cnt += 1  # 먹은 물고기 수 증가
    start = (x, y)  # 상어의 위치 업데이트
    if cnt == large:  # 먹은 물고기 수가 상어의 크기와 같다면
        cnt = 0  # 먹은 물고기 수 초기화
        large += 1  # 상어의 크기 증가

print(res)  # 총 걸린 시간 출력

문제를 풀고 나니, 물고기를 찾으면 그 for문 까지만 확인하고 멈춘다면 시간과 메모리를 많이 단축했을 것 같다는 생각이 들었다. 물론, 다시 하지는 않았지만.............