[Python] 백준 16236. 아기상어 (BFS)

문제 링크

문제 풀이

1. 상어 위치에서부터 bfs로 잡아먹을 수 있는 가장 가까운 물고기를 찾는다
   • 이때, 문제 조건대로 가까운 물고기가 많으면, 가장 위에, 다음에 가장 왼쪽에 있는 물고기를 찾는다
2. 그런 물고기가 존재하지 않으면,
   • 엄마 상어에게 도움 요청함
   • 즉, time 출력하고 프로그램 종료
3. 물고기가 존재하면,
   • 그 물고기가 있는 칸으로 가서 먹고,
   • 이동한 시간을 업데이트하고
   • 다시 1번부터 반복

코드

import itertools
import sys
from collections import deque

input = sys.stdin.readline

dx = [-1, 0, 1, 0]  # 상좌하우
dy = [0, -1, 0, 1]


def bfs_eatable_fish(s_x, s_y):
    q = deque([[0, s_x, s_y]])
    visited = [[False for _ in range(N)] for _ in range(N)]
    eatable_fish = []

    while q:
        dist, x, y = q.popleft()
        for i, j in zip(dx, dy):
            n_x, n_y = x + i, y + j
            if (
                0 <= n_x < N
                and 0 <= n_y < N
                and not visited[n_x][n_y]
                and graph[n_x][n_y] <= shark_size
            ):
                visited[n_x][n_y] = True
                q.append([dist + 1, n_x, n_y])
                if 0 < graph[n_x][n_y] < shark_size:
                    eatable_fish.append([dist + 1, n_x, n_y])

    eatable_fish.sort(
        key=lambda x: (x[0], x[1], x[2])
    )  # 먹을 수 있는 물고기들을 거리순, x, y 순으로 정렬한 리스트

    return eatable_fish


if __name__ == "__main__":
    N = int(input())
    graph = [list(map(int, input().strip().split())) for _ in range(N)]
    shark_size, eat_cnt, time = 2, 0, 0
    s_x, s_y = 0, 0  # shark's position

    for x, y in itertools.product(range(N), range(N)):  #상어 위치 찾기
        if graph[x][y]:
            s_x, s_y = x, y
            graph[x][y] = 0
            break

    while True:
        next_fish_list = bfs_eatable_fish(s_x, s_y) #상어 위치에서부터 bfs로 잡아먹을 수 있는 가장 가까운 물고기를 찾는다
        if not next_fish_list: 
            print(time)
            exit(0)

        dist, x, y = next_fish_list[0]
        s_x, s_y = x, y #상어 위치 업데이트
        time += dist
        eat_cnt += 1    #물고기 먹음
        graph[x][y] = 0
        if eat_cnt == shark_size:
            eat_cnt = 0
            shark_size += 1