[백준-16236] 아기 상어 / Python

16236번: 아기 상어

import sys
from collections import deque

input = sys.stdin.readline

n = int(input())
graph = []
for i in range(n):
    li = list(map(int, input().split()))
    graph.append(li)
    for j in range(n):
        if li[j] == 9:
            graph[i][j] = 2
            start = [i, j]

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


def bfs(i, j):
    visited = [[0] * n for _ in range(n)]
    visited[i][j] = 1
    distance = [[0] * n for _ in range(n)]
    eat = []
    queue = deque()
    queue.append((i, j))
    while queue:
        x, y = queue.popleft()
        for k in range(4):
            nx = x + dx[k]
            ny = y + dy[k]
            if 0 <= nx < n and 0 <= ny < n and visited[nx][ny] == 0:
                if graph[nx][ny] <= graph[i][j]:
                    visited[nx][ny] = 1
                    distance[nx][ny] = distance[x][y] + 1
                    queue.append((nx, ny))
                if graph[nx][ny] < graph[i][j] and graph[nx][ny] != 0:
                    eat.append((nx, ny, distance[nx][ny]))
    if not eat:
        return -1, -1, -1
    eat.sort(key=lambda x: (x[2], x[0], x[1]))
    return eat[0][0], eat[0][1], eat[0][2]


exp = 0
answer = 0
while True:
    i, j = start[0], start[1]
    ex, ey, dist = bfs(i, j)
    if ex == -1:
        break
    graph[ex][ey] = graph[i][j]
    graph[i][j] = 0
    start = [ex, ey]
    exp += 1
    if exp == graph[ex][ey]:
        exp = 0
        graph[ex][ey] += 1
    answer += dist

print(answer)

 

해설

bfs 문제다.

먹을 수 있는 물고기의 위치와 거리를 eat에 저장한다.

lambda를 사용해 거리, x, y를 기준으로 정렬하여 첫 번째 인덱스의 데이터를 리턴한다.

새로운 시작 위치는 해당 인덱스의 x, y 값이고, 아기 상어의 위치를 새로운 시작 위치로 옮긴다. 이 때, 기존 아기 상어의 위치는 0으로 변경한다.

새로운 시작 위치(물고기를 먹었을 때)가 갱신될 때마다 exp를 하나씩 늘려주며 아기 상어의 크기와 비교한다.

아기 상어의 크기와 exp가 같다면 exp는 0으로, 아기 상어의 크기는 기존의 크기에 1을 더해준다.

answer에 리턴해준 먹을 수 있는 물고기까지의 거리를 더해준다.