300. 상어 중학교

백준

1. 구현

from collections import deque
import sys
input = sys.stdin.readline
dx = [-1, 1, 0, 0]
dy = [0, 0, -1, 1]

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

def rotate(board):
    temp = [[0] * n for _ in range(n)]
    for j in range(n):
        for i in range(n):
            temp[n-j-1][i] = board[i][j]

    return temp

def bfs(x, y, p):
    q.append([x, y])
    temp = board[x][y]
    check[x][y] = p
    cnt, r = 1, 0
    while q:
        x, y = q.pop()
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < n and 0 <= ny < n:
                if board[nx][ny] == temp and check[nx][ny] == 0:
                    check[nx][ny] = p
                    cnt += 1
                    q.append([nx, ny])
                elif board[nx][ny] == 0 and p not in check[nx][ny]:
                    check[nx][ny].append(p)
                    cnt += 1
                    r += 1
                    q.append([nx, ny])
    return cnt, r


def fall(x, y):
    flag = 0
    for i in range(x+1, n):
        nx = i
        if board[i][y] > -2:
            flag = 1
            break
    if flag:
        board[nx-1][y] = board[x][y]
    else:
        board[nx][y] = board[x][y]
    board[x][y] = -2




ans = 0
while True:
    q = deque()
    check = [[0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            if board[i][j] == 0:
                check[i][j] = [] #무지개는 빈 리스트 생성
    p, b = 1, []
    for i in range(n):
        for j in range(n):
            if board[i][j] > 0 and check[i][j] == 0:
                cnt, r = bfs(i, j, p)
                if cnt > 1:
                    b.append([cnt, r, i, j, p])
                p += 1
    #print(check)
    if not b:
        break

    b = sorted(b)

    cnt = 0
    for i in range(n):
        for j in range(n):
            if board[i][j] > 0 and check[i][j] == b[-1][-1]:
                board[i][j] = -2
                cnt += 1
            elif board[i][j] == 0 and b[-1][-1] in check[i][j]:
                board[i][j] = -2
                cnt += 1
    ans += cnt ** 2

    for i in range(n-2, -1, -1):
        for j in range(n):
            if board[i][j] >= 0 and board[i+1][j] == -2:
                fall(i, j)

    board = rotate(board)
    #board = list(zip(*board))[::-1] 
    #board = [list(s) for s in board]

    for i in range(n-2, -1, -1):
        for j in range(n):
            if board[i][j] >= 0 and board[i+1][j] == -2:
                fall(i, j)

print(ans)