[acmicpc] 16236. 아기상어(python)

import sys
from collections import deque
input=sys.stdin.readline

N=int(input())
arr=[list(map(int,input().split())) for _ in range(N)]
time=0
babySize=2
dirs=[(1,0),(0,1),(-1,0),(0,-1)]
for r in range(N):
    for c in range(N):
        if arr[r][c]==9: baby=(r,c,0)
answer=0
cnt=0
while True:
    q=deque()
    q.append(baby)
    visit=[[0]*N for _ in range(N)]
    visit[baby[0]][baby[1]]=1
    remove=[]
    while q:
        r,c,time=q.popleft()
        if 1<=arr[r][c]<babySize:
            remove.append((time,r,c))
        for i in range(4):
            nr,nc=r+dirs[i][0],c+dirs[i][1]
            if 0>nr or nr>=N or nc<0 or nc>=N: continue
            if arr[nr][nc]>babySize: continue
            if visit[nr][nc]==1: continue
            q.append((nr,nc,time+1))
            visit[nr][nc]=1
    if len(remove) == 0: break
    remove.sort()
    answer+=remove[0][0]
    arr[remove[0][1]][remove[0][2]]=0
    arr[baby[0]][baby[1]] = 0
    baby=(remove[0][1],remove[0][2],0)
    cnt+=1
    if cnt==babySize:
        cnt=0
        babySize+=1
print(answer)

---

https://www.acmicpc.net/problem/16236

16236번: 아기 상어