BOJ - 21609 - 상어 중학교

import sys
from collections import deque

__BLANK = -2
__BLACK = -1
__RAINBOW = 0
dx = [-1, 0, 1, 0]
dy = [0, 1, 0, -1]

def bfs(x, y, color, SHARK, visit, N):
    visit[x][y] = True
    dq = deque()
    dq.append((x, y))
    
    block_cnt = 1
    rainbow_cnt = 0
    block_list, rainbow_list = [(x,y)], []

    while dq:
        x, y = dq.popleft()
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            
            if 0 <= nx < N and 0 <= ny < N and not visit[nx][ny]:
                if SHARK[nx][ny] == color:
                    visit[nx][ny] = True
                    dq.append((nx, ny))
                    block_cnt += 1
                    block_list.append((nx, ny))
                
                elif SHARK[nx][ny] == __RAINBOW:
                    visit[nx][ny] = True
                    dq.append((nx, ny))
                    block_cnt += 1
                    rainbow_cnt += 1
                    rainbow_list.append((nx, ny))

    for row, col in rainbow_list:
        visit[row][col] = False

    return block_cnt, rainbow_cnt, block_list + rainbow_list

def remove(SHARK, Block):
    for x, y in Block:
        SHARK[x][y] = __BLANK

def rotation90(SHARK, N):
    result = [[0 for _ in range(N)] for _ in range(N)]
    for row in range(N):
        for col in range(N):
            result[N - 1 - col][row] = SHARK[row][col]

    return result
    
def gravity(SHARK, N):
    for row in range(N-2, -1, -1):
        for col in range(N):
            if SHARK[row][col] > __BLACK:
                current = row
                for down in range(row+1, N):
                    if SHARK[down][col] == __BLANK:
                        SHARK[down][col] = SHARK[current][col]
                        SHARK[current][col] = __BLANK
                        current += 1
                    else: break

def solution():
    N, M = map(int, sys.stdin.readline().split())
    SHARK = [list(map(int, sys.stdin.readline().split())) for _ in range(N)]
    
    score = 0
    while True:

        visit = [[False for _ in range(N)] for _ in range(N)]

        isBlocks = []
        for row in range(N):
            for col in range(N):
                if SHARK[row][col] > __RAINBOW and not visit[row][col]:
                    block_count, rainbow_cnt, block_list = bfs(row, col, SHARK[row][col], SHARK, visit, N)
                    if block_count >= 2:
                        isBlocks.append((block_count, rainbow_cnt, block_list))

        if not isBlocks:
            break

        isBlocks = sorted(isBlocks, reverse=True)
        score += isBlocks[0][0] ** 2
        remove(SHARK, isBlocks[0][2])
        gravity(SHARK, N)
        SHARK = rotation90(SHARK, N)
        gravity(SHARK, N)

    print(score)
solution()

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