[백준] 16236 아기 상어 with 파이썬

난 상어가 무섭다

 

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

INF = 1e9
N = int(input())
array = [list(map(int, input().split())) for _ in range(N)]

# 필요한 변수 초기화
shark_size = 2
eated_number = 0        # 상어 크기 증가를 위해 필요한 변수
time = 0

dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

# 초기 상어 위치 파악하기, 먹을 물고기 파악하기
for i in range(N):
    for j in range(N):
        if array[i][j] == 9:
            shark_x, shark_y = i, j
            array[i][j] = 0
def bfs(start_x, start_y, shark_size):  # 상어크기를 고려한 bfs
    queue = deque()
    queue.append((start_x, start_y))
    range_array[start_x][start_y] = 0

    while queue:
        x, y = queue.popleft()
        for i in range(4):
            nx, ny = x + dx[i], y + dy[i]
            if (0 <= nx < N) and (0 <= ny < N) and (array[nx][ny] <= shark_size) and (range_array[nx][ny] == 0):
                queue.append((nx, ny))
                range_array[nx][ny] = range_array[x][y] + 1

while(1):
    eatable_fishs = []
    range_array = [[0 for _ in range(N)] for _ in range(N)]

    # 먹을 물고기 탐색
    for i in range(N):
        for j in range(N):
            if (0 < array[i][j]) and (array[i][j] < shark_size):
                eatable_fishs.append([i, j])

    # # 먹을 물고기가 없으면 break
    if len(eatable_fishs) == 0:
        break

    # 먹을 물고기가 있으면 bfs로 거리 맵 만들기
    bfs(shark_x, shark_y, shark_size)

    # 거리 정보 더하기
    can_move = 0
    for i in range(len(eatable_fishs)):
        if range_array[eatable_fishs[i][0]][eatable_fishs[i][1]] != 0:
            eatable_fishs[i].append(range_array[eatable_fishs[i][0]][eatable_fishs[i][1]])
        else:   # 거리가 0이라면 도달할 수 없다는 뜻이므로 거리 무한으로 설정하기
            eatable_fishs[i].append(INF)
    
    # 가장 우선순위 높은 물고기로 만들기
    eatable_fishs = sorted(eatable_fishs, key = lambda x : (x[2], x[0], x[1]))

    # 가장 우선순위 높은 물고기로 이동
    shark_x, shark_y = eatable_fishs[0][0], eatable_fishs[0][1]
    array[shark_x][shark_y] = 0
    if eatable_fishs[0][2] != INF:   # 거리가 무한이 아닐때
        time += eatable_fishs[0][2]
    else:
        break
    eated_number += 1
    eatable_fishs = []

    # 먹은 개수가 상어 크기만해진다면 상어 크기 증가
    if eated_number == shark_size:
        shark_size += 1
        eated_number = 0

print(time)

# # 확인용
# for i in range(N):
#     for j in range(N):
#         print(array[i][j], end = ' ')
#     print()
# print()