[acmicpc] 21609. 상어 중학교(python)

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

N,M=map(int,input().split())
arr=[list(map(int,input().split())) for _ in range(N)]
dirs=[(1,0),(0,1),(-1,0),(0,-1)]
score=0

def bfs(row,col,color):
    global removes, cmax
    visit = [[0] * N for _ in range(N)]
    q=deque([(row,col)])
    colors=[(row,col)]
    visit[row][col]=1
    remove=[(row,col)]
    cnt=0
    mcnt=1 # 일반 블록 개수
    while q:
        r,c=q.popleft()
        cnt+=1
        for i in range(4):
            nr,nc=r+dirs[i][0],c+dirs[i][1]
            if nr<0 or nr>=N or nc<0 or nc>=N : continue
            if (arr[nr][nc]==color or arr[nr][nc]==0) and visit[nr][nc]==0:
                if arr[nr][nc]==color:
                    colors.append((nr,nc))
                    mcnt+=1
                remove.append((nr,nc))
                visit[nr][nc]=1
                q.append((nr,nc))
    for r,c in colors:
        cvisit[r][c]=1
    if cnt >= 2:
        removes.append([cnt,cnt-mcnt,remove])
def gravity():
    for c in range(N):
        for _ in range(N):
            for r in range(1, N):
                if arr[r][c] == -2 and arr[r - 1][c] != -1:
                    arr[r][c] = arr[r - 1][c]
                    arr[r - 1][c] = -2
while True:
    cvisit=[[0] * N for _ in range(N)]
    removes=[]
    cmax=0
    # 블록 찾기
    for r in range(N):
        for c in range(N):
            if arr[r][c]>0 and cvisit[r][c]==0:
                bfs(r,c,arr[r][c])
    if len(removes)==0: break
    # 블록 제거
    removes.sort(reverse=True)
    score += pow(removes[0][0], 2)
    for r,c in removes[0][2]:
        arr[r][c]=-2
    # 중력 작용
    gravity()
    # 반시계 90도 회전
    arr=[list(l) for l in list(reversed(list(zip(*arr))))]
    # 중력 작용
    gravity()
print(score)

---

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

21609번: 상어 중학교