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

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

21609번: 상어 중학교

 

BFS 를 활용하는 시뮬레이션 문제이다.

구현할 때 모든 조건을 꼼꼼히 살펴야하고, 작은 사소한 실수 하나로도 문제를 틀리기 때문에 난이도가 높다.

N, M = map(int, input().split())
board = [list(map(int, input().split())) for _ in range(N)]
score = 0
while True:
    x, y, size = find_block_group()
    if size < 2:
        break

    delete_block_group(x, y)
    score += (size * size)
    
    set_gravity()
    
    rotate()
    
    set_gravity()


print(score)

 

이런 복잡한 문제는 큰 틀부터 잡은 다음, 디테일을 구현해나가는 것이 구현하기 편하다.

먼저 find_block_group() 함수부터 구현했다.

 

def find_block_group():
    _x, _y, max_size = None, None, 0
    max_rainbow_block = 0
    visited = [[False for _ in range(N)] for _ in range(N)]
    for i in range(N):
        for j in range(N):
            if board[i][j] in (0, -1) or visited[i][j]:
                continue
            d = deque()
            d.append((i, j))
            color = board[i][j]
            size = 1
            visited[i][j] = True
            rainbow_block = 0
            while d:
                now_i, now_j = d.popleft()
                for k in range(4):
                    next_i, next_j = now_i + di[k], now_j + dj[k]
                    if 0 <= next_i < N and 0 <= next_j < N:
                        if not visited[next_i][next_j] and board[next_i][next_j] in (color, 0):
                            if board[next_i][next_j] == 0:
                                rainbow_block += 1
                            d.append((next_i, next_j))
                            visited[next_i][next_j] = True
                            size += 1

            if size > max_size:
                max_size = size
                _x, _y = i, j
            elif size == max_size:
                if rainbow_block > max_rainbow_block:
                    max_rainbow_block = rainbow_block
                    _x, _y = i, j
                elif rainbow_block == max_rainbow_block:
                    _x, _y = i, j # 나중에 탐색한 게 기준 블록의 행/열이 더 크다.

            # 무지개 블록의 방문 초기화
            for _i in range(N):
                for _j in range(N):
                    if board[_i][_j] == 0:
                        visited[_i][_j] = False

    return _x, _y, max_size

 

이 코드는 내가 맞왜틀을 했던 코드이다.

이 코드에서 단 한줄을 추가하지 않아서 제출시 0%에서 틀렸다.

if size > max_size:
    max_size = size
    _x, _y = i, j
elif size == max_size:
    if rainbow_block > max_rainbow_block:
        max_rainbow_block = rainbow_block
        _x, _y = i, j

 

바로 이 부분이 잘못되었다.

블록 그룹에서 우선순위를 판단할 때는 블록 그룹의 크기, 블록 그룹 내 무지개 블록의 개수가 사용된다.

이때 블록 그룹의 크기를 가지고 판단을 한 상황에서 "무지개 블록의 개수가 바뀌지 않고 있다."

그래서 아래와 같이 한 줄의 코드를 추가해주면 된다.

if size > max_size:
    max_size = size
    max_rainbow_block = rainbow_block
    _x, _y = i, j
elif size == max_size:
    if rainbow_block > max_rainbow_block:
        max_rainbow_block = rainbow_block
        _x, _y = i, j

 

이 함수는 조건에 맞는 블록 그룹의 크기, 행, 열 좌표를 반환한다.

만약 이 함수가 반환한 블록 그룹의 크기가 1 이하라면 지울 수 있는 블록 그룹이 없으므로 반복문을 빠져나간다.

 

이제 지울 수 있는 블록 그룹이 있을 때, 블록 그룹을 지우는 함수를 작성하자.

def delete_block_group(x, y):
    color = board[x][y]
    d = deque()
    d.append((x, y))
    board[x][y] = -2
    while d:
        now_i, now_j = d.popleft()
        for k in range(4):
            next_i, next_j = now_i + di[k], now_j + dj[k]
            if 0 <= next_i < N and 0 <= next_j < N:
                if board[next_i][next_j] in (color, 0):
                    d.append((next_i, next_j))
                    board[next_i][next_j] = -2

 

삭제할 블록 그룹의 기준 블록 좌표를 입력하면, 그 좌표에서부터 다시 BFS를 돌면서 모든 블럭을 지우는 코드이다.

 

이 함수를 구현할 때도 한가지 실수한 부분이 있었다.

BFS 를 구현할 때 방문체크 (이 함수에서는 블록을 지웠다는 표시로 -2를 하면서 체크) 는 큐에 넣으면서 체크해야한다.

큐에서 꺼내면서 체크하면 안된다.

 

하지만 난 이 당연한 사실을 알면서도 큐에서 꺼내면서 체크를 하는 실수를 저지르고 말았다..

visit 배열을 쓰지 않고 구현해서 낯설었던 모양이다.

 

다음으로 중력을 적용하는 함수를 구현했다.

def set_gravity():
    for j in range(N):
        floor = N-1
        for i in range(N-1, -1, -1):
            if board[i][j] == -1:
                floor = i-1
            elif board[i][j] >= 0:
                if i != floor:
                    board[floor][j] = board[i][j]
                    board[i][j] = -2
                floor -= 1

 

아래에서부터 위로 훑으면서 내릴 수 있는 블럭을 만날 때마다, 바닥 좌표로 해당 블럭을 옮겨주고, 바닥 좌표를 해당 블럭 위로 바꿔주었다.

 

마지막으로 반시계방향 90도 회전 코드를 작성하였다.

def rotate():
    copied_board = [[board[i][j] for j in range(N)] for i in range(N)]
    for i in range(N):
        for j in range(N):
            board[N-1-j][i] = copied_board[i][j]

 

프로그래밍 언어에서 2차원 배열을 배울 때 많이 하던 그 작업이다.

 

전체 코드는 아래와 같다.

from collections import deque
di = [-1, 0, 0, 1]
dj = [0, -1, 1, 0]


def find_block_group():
    _x, _y, max_size = None, None, 0
    max_rainbow_block = 0
    visited = [[False for _ in range(N)] for _ in range(N)]
    for i in range(N):
        for j in range(N):
            if board[i][j] in (0, -1, -2) or visited[i][j]:
                continue
            d = deque()
            d.append((i, j))
            color = board[i][j]
            size = 1
            visited[i][j] = True
            rainbow_block = 0
            while d:
                now_i, now_j = d.popleft()
                for k in range(4):
                    next_i, next_j = now_i + di[k], now_j + dj[k]
                    if 0 <= next_i < N and 0 <= next_j < N:
                        if not visited[next_i][next_j] and board[next_i][next_j] in (color, 0):
                            if board[next_i][next_j] == 0:
                                rainbow_block += 1
                            d.append((next_i, next_j))
                            visited[next_i][next_j] = True
                            size += 1

            if size > max_size:
                max_size = size
                max_rainbow_block = rainbow_block
                _x, _y = i, j
            elif size == max_size:
                if rainbow_block > max_rainbow_block:
                    max_rainbow_block = rainbow_block
                    _x, _y = i, j
                elif rainbow_block == max_rainbow_block:
                    _x, _y = i, j # 나중에 탐색한 게 기준 블록의 행/열이 더 크다.

            # 무지개 블록의 방문 초기화
            for _i in range(N):
                for _j in range(N):
                    if board[_i][_j] == 0:
                        visited[_i][_j] = False

    return _x, _y, max_size


def delete_block_group(x, y):
    color = board[x][y]
    d = deque()
    d.append((x, y))
    board[x][y] = -2
    while d:
        now_i, now_j = d.popleft()
        for k in range(4):
            next_i, next_j = now_i + di[k], now_j + dj[k]
            if 0 <= next_i < N and 0 <= next_j < N:
                if board[next_i][next_j] in (color, 0):
                    d.append((next_i, next_j))
                    board[next_i][next_j] = -2


def set_gravity():
    for j in range(N):
        floor = N-1
        for i in range(N-1, -1, -1):
            if board[i][j] == -1:
                floor = i-1
            elif board[i][j] >= 0:
                if i != floor:
                    board[floor][j] = board[i][j]
                    board[i][j] = -2
                floor -= 1


def rotate():
    copied_board = [[board[i][j] for j in range(N)] for i in range(N)]
    for i in range(N):
        for j in range(N):
            board[N-1-j][i] = copied_board[i][j]


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

score = 0
while True:
    x, y, size = find_block_group()
    if size < 2:
        break

    delete_block_group(x, y)

    score += (size * size)

    set_gravity()

    rotate()

    set_gravity()

print(score)