Jae-Kyung Cho LLM Developers who was a Robotics engineer

Baekjoon-16236-아기 상어

from collections import deque

N = int(input())
board = [list(map(int,input().split())) for _ in range(N)]

fishes = {k:[] for k in range(1,7)}
shark = []
for i in range(N):
    for j in range(N):
        if board[i][j]==9:
            shark = [i,j,2,0] # x,y,size,fish count
            board[i][j] = 0
        if 1<=board[i][j]<=6:
            fishes[board[i][j]].append([i,j])

result = 0
while True:
    q = deque([[0,shark[0],shark[1]]])
    visited = [[-1]*N for _ in range(N)]
    visited[shark[0]][shark[1]] = 1
    bfs_result = [401,-1,-1]
    while q:
        d,x,y = q.popleft()
        if d > bfs_result[0]:
            continue
        for dx, dy in [[-1,0],[0,-1],[0,1],[1,0]]:
            nx,ny = x+dx, y+dy
            if 0<=nx<N and 0<=ny<N and visited[nx][ny]==-1 and board[nx][ny] <= shark[2]:
                visited[nx][ny] = d+1
                q.append([d+1,nx,ny])
                if 0 < board[nx][ny] < shark[2]:
                    if bfs_result[0] > d+1:
                        bfs_result = [d+1,nx,ny]
                    elif bfs_result[0] == d+1:
                        if nx < bfs_result[1]:
                            bfs_result = [d+1,nx,ny]
                        elif nx==bfs_result[1] and ny < bfs_result[2]:
                            bfs_result = [d+1,nx,ny]
    # print('----------')
    # print(*board,sep='\n')
    if bfs_result[0]!=401:
        d,nx,ny = bfs_result
        result += d
        shark[0], shark[1] = nx,ny
        shark[3] += 1
        if shark[3] == shark[2]:
            shark[2] += 1
            shark[3] = 0
        board[nx][ny] = 0
    else:
        break

print(result)
Comments