[Baekjoon] 16236 아기상어 python

🏷 문제

💡 코드

from collections import deque

n = int(input())
space = []
for i in range(n):
    space.append(list(map(int, input().split())))

# 상하좌우로 이동한다고 가정
case = [[-1, 0],[0, -1], [0, 1], [1, 0]]

# 상어 크기
shark_size = 2

def find(x,y):
    visited = [[0 for _ in range(n)] for _ in range(n)]
    visited[x][y] = 1
    q = deque()
    q.append([x, y])
    res = []
    while q:
        x, y = q.popleft()
        # 상하좌우 탐색
        for i in range(4):
            newx = x + case[i][0]
            newy = y + case[i][1]
            if 0 <= newx < n and 0 <= newy < n:
                # 이동 가능한지 확인(물고기 크기, 이전에 방문했는지에 대한 여부)
                if space[newx][newy] <= shark_size and visited[newx][newy] == 0:
                    # 먹을 수 있는지 확인
                    if space[newx][newy] < shark_size and space[newx][newy] != 0:
                        res.append([newx, newy, visited[x][y]])
                    else:
                        # 이동한 위치를 넣어줘야 그 다음부턴 그 위치를 기준으로 계산할 수 있음
                        # ex_거리가 1일때, 2일때, 3일때 ...
                        q.append([newx,newy])
                        # 탐색하는 데 걸린 시간을 업데이트
                        visited[newx][newy] = visited[x][y] + 1
    # 상어가 물고기를 하나도 먹지 못한 경우
    if len(res) == 0:
        return -1, -1, -1
    else:
        # x[0] = x좌표(상하), x[1] = y좌표(좌우), x[2] = 물고기를 먹기까지 걸리는 시간
        res = sorted(res, key=lambda x: (x[2], x[0], x[1]))
        return res[0]

# 상어가 물고기를 먹는데 걸리는 총 시간
time = 0
# 상어가 물고기를 먹은 횟수
shark_cnt = 0

def eat_fish():
    global space, shark_cnt, shark_size, time
    for i in range(n):
        for j in range(n):
            if space[i][j] == 9:
                sharkx, sharky, sec = find(i,j)
                if sharkx == -1 and sharky == -1 and sec == -1:
                    return False
                else:
                    shark_cnt += 1
                    time += sec
                    # 상어 크기과 먹은 물고기 수가 같으면
                    if shark_size == shark_cnt:
                        shark_size += 1
                        shark_cnt = 0
                    # 현재 상어의 위치를 0으로 만들기
                    space[i][j] = 0
                    # 상어가 이동할 위치를 9로 만들기
                    space[sharkx][sharky] = 9
                    return True

while True:
    if eat_fish():
        continue
    else:
        break

print(time)

🔑