[BOJ] 21609 - 상어 중학교 (Python)

 

 

아이디어

 

격자에 중력과 회전을 적용하며 BFS로 블록 그룹을 찾아 점수 합을 계산한다.

 

 

풀이

 

import sys

input = sys.stdin.readline

def solution():
    delta = [(-1, 0), (0, 1), (1, 0), (0, -1)]

    def gravity():
        for c in range(N):
            end = N-1
            for r in range(N-1, -1, -1):
                if board[r][c] >= 0:
                    board[end][c] = board[r][c]
                    if end != r:
                        board[r][c] = -2
                    end -= 1
                elif board[r][c] == -1:
                    end = r-1

    def rotate():
        nonlocal board
        board = list(map(list, zip(*board)))
        for r in range(N//2):
            board[r], board[-(r+1)] = board[-(r+1)], board[r]

    N, M = map(int, input().split())
    board = [list(map(int, input().split())) for _ in range(N)]
    res = 0
    while 1:
        visited = [[0]*N for _ in range(N)]
        groups = []
        idx = 0
        for i in range(N):
            for j in range(N):
                if visited[i][j] == 0 and board[i][j] > 0:
                    visited[i][j] = -1
                    queue = [(i, j)]
                    idx += 1
                    cnt = 1
                    rainbow = 0
                    while queue:
                        r, c = queue.pop(0)
                        for dr, dc in delta:
                            nr, nc = r+dr, c+dc
                            if N > nr >= 0 and N > nc >= 0:
                                if visited[nr][nc] not in [-1, idx]:
                                    if board[nr][nc] == board[i][j]:
                                        visited[nr][nc] = -1
                                        queue.append((nr, nc))
                                        cnt += 1
                                    elif board[nr][nc] == 0:
                                        visited[nr][nc] = idx
                                        queue.append((nr, nc))
                                        cnt += 1
                                        rainbow += 1
                    if cnt > 1:
                        groups.append((cnt, rainbow, i, j))
        if groups:
            groups.sort(key=lambda x: (-x[0], -x[1], -x[2], -x[3]))
            i, j = groups[0][2], groups[0][3]
            color = board[i][j]
            board[i][j] = -2
            queue = [(i, j)]
            cnt = 1
            while queue:
                r, c = queue.pop(0)
                for dr, dc in delta:
                    nr, nc = r+dr, c+dc
                    if N > nr >= 0 and N > nc >= 0:
                        if board[nr][nc] in [0, color]:
                            board[nr][nc] = -2
                            cnt += 1
                            queue.append((nr, nc))
            res += cnt**2
        else:
            break
        gravity()
        rotate()
        gravity()
    print(res)

solution()