백준 #16236 아기상어 (파이썬)

문제 링크

• 먹을 수 있는 물고기가 1마리보다 많다면, 거리가 가장 가까운 물고기를 먹으러 간다.
  • 거리는 아기 상어가 있는 칸에서 물고기가 있는 칸으로 이동할 때, 지나야하는 칸의 개수의 최솟값이다.
  • 거리가 가까운 물고기가 많다면, 가장 위에 있는 물고기, 그러한 물고기가 여러마리라면, 가장 왼쪽에 있는 물고기를 먹는다.

이게 문제의 핵심이고, 이미 짐작하듯이 bfs를 이용한 최단거리 찾기 관련 문제다.

하나의 물고기를 먹으려 할 때마다 전체 space에 대하여 bfs를 진행해서, 먹을 수 있는 물고기를 다 정리하고, 그 중 가장 조건에 맞는 물고기 단 하나를 먹는 행위를 무한 반복하는 것이 알고리즘이다.

먹을 수 있는 물고기가 하나도 없을 때까지(if not food) 반복한다.

매우 시간을 오래 잡아먹을 것 같이 보이지만, 196ms밖에 걸리지 않았다.

사실 천천히 들여다보면, dfs와 bfs는 적어도 했던 짓을 또하진 않는다.
그래서 전체 arr의 크기(여기서는 N*N)의 시간복잡도만을 가질 뿐이다.

비슷한 arr에 대해 여러번 탐색을 진행하더라도 그렇게 시간이 크게 소모되지 않으니 크게 걱정하지 말자.

가독성을 챙기려 굉장히 노력한 코드다.

import sys

N = int(sys.stdin.readline().rstrip())

space = []
for _ in range(N):
    space.append(list(map(int, sys.stdin.readline().rstrip().split())))

# 처음 아기 상어의 크기는 2
# 1 두개를 먹으면 3이 된다. 

# bfs로 자기가 먹을 수 있는 물고기를 찾는다. 
# 물고기 위치와 거리를 저장해두고, 거리 -> y좌표 -> x좌표 순으로 정렬해준다. 
# 먹을 수 있는 거 다 먹고 이동시간에 더한다. 
# 반복.

class BabyShark:
    def __init__(self, y, x, size):
        self.y = y
        self.x = x
        self.size = size
        self.cnt = 0
        self.time = 0
    
    def move(self, distance, y, x):
        self.time += distance
        self.y = y
        self.x = x

    def eat(self):
        self.cnt += 1
        if self.cnt == self.size:
            self.size += 1
            self.cnt = 0


DIRECTIONS = [(-1,0), (1,0), (0,-1), (0,1)]

if __name__ == '__main__':

    for i in range(N):
        for j in range(N):
            if space[i][j] == 9:
                space[i][j] = 0
                babyShark = BabyShark(i, j, 2)
                break

    while True:
        visited = [[False] * N for _ in range(N)]
        food = []

        q = [(babyShark.y, babyShark.x)]
        distance = 0
        while q:
            q_cpy = q[:]
            q = []
            distance += 1
            for y, x in q_cpy:
                for dy, dx in DIRECTIONS:
                    if not (0 <= y+dy < N and 0 <= x+dx < N):
                        continue
                    if visited[y+dy][x+dx]:
                        continue
                    if space[y+dy][x+dx] > babyShark.size:
                        continue
                    q.append((y+dy, x+dx))
                    visited[y+dy][x+dx] = True
                    if 0 < space[y+dy][x+dx] < babyShark.size:
                        food.append((distance, y+dy, x+dx))

        if not food:
            print(babyShark.time)
            exit()

        food.sort()
        best_food = food[0]
        best_food_distance, best_food_y, best_food_x = best_food
        babyShark.move(best_food_distance, best_food_y, best_food_x)
        babyShark.eat()
        space[best_food_y][best_food_x] = 0