[코드트리] 색깔 폭탄 - python

문제 링크: https://www.codetree.ai/frequent-problems/colored-bomb/description

코드트리

bfs를 활용한 시뮬레이션 문제. 굉장히 복잡한 단계들을 거쳐야하는 문제여서 보자마자 한숨이 나왔다..

 

bfs를 통해 폭탄 묶음들을 구한 다음 조건에 맞게 비교해서 터뜨려야할 폭탄 묶음을 구해야한다. 조건 중 기준점을 구해야하는데 빨간 폭탄은 기준점에서 제외해야하는데 이를 체크해주지 않아서 디버깅하는데 꽤 많은 시간을 사용하였다. 조건에 대한 예외를 생각해보고 체크해보는 것이 중요한 것 같다.

 

걸린 시간: 1시간 54분 28초

# 문제 링크: https://www.codetree.ai/frequent-problems/colored-bomb/description
from collections import deque
import sys
sys.stdin = open('input.txt')

n, m = tuple(map(int, input().split()))

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

RED = 0
STONE = -1
EMPTY = -2

def inRange(y, x):
    return 0 <= y < n and 0 <= x < n

def get_bundle(y, x, visited):
    dys = [1, 0, -1, 0]
    dxs = [0, 1, 0, -1]

    temp = [
        [False] * n
        for _ in range(n)
    ]

    temp[y][x] = True
    color = board[y][x]
    bundle = []
    count_of_red = 0
    q = deque([(y, x)])


    while q:
        y, x = q.popleft()
        # print(y, x)

        bundle.append((y, x))
        if board[y][x] == 0:
            count_of_red += 1

        for dy, dx in zip(dys, dxs):
            ny = y + dy
            nx = x + dx
            # 돌이 아니고 빈 곳도 아니고 빨간돌이거나 아직 방문하지 않은 색이 같은 폭탄인 경우
            if inRange(ny, nx) and board[ny][nx] != STONE and board[ny][nx] != EMPTY and not temp[ny][nx] and (board[ny][nx] == RED or board[ny][nx] == color):
                q.append((ny, nx))
                temp[ny][nx] = True

    # 방문한 묶음들은 표시, 단 빨간 폭탄은 제외한다.
    for i in range(n):
        for j in range(n):
            if board[i][j] != RED and temp[i][j]:
                visited[i][j] = temp[i][j]

    return (bundle, count_of_red)


def get_bombs():
    bombs = []
    visited = [
        [False] * n
        for _ in range(n)
    ]

    for i in range(n):
        for j in range(n):
            # 돌이 아니고 빨간돌이 아니고 빈 곳도 아니고 방문하지 않은 경우
            if board[i][j] != STONE and board[i][j] != RED and board[i][j] != EMPTY and not visited[i][j]:
                bundle, count_of_red = get_bundle(i, j, visited)

                if len(bundle) > 1:
                    bombs.append((count_of_red, bundle))
    return bombs

def get_point(pos):
    t_y = -1
    t_x = -1

    for y, x in pos:
        if board[y][x] == RED:
            continue

        if t_y < y:
            t_y = y
            t_x = x

        if t_y == y:
            if t_x > x:
                t_x = x

    return (t_y, t_x)


def get_target(bombs):
    max_of_bomb = 0
    min_of_red = 0
    target = []

    for count_of_red, bomb in bombs:
        # 폭탄의 개수가 가장 큰 것 선택
        if len(bomb) > max_of_bomb:
            max_of_bomb = len(bomb)
            min_of_red = count_of_red
            target = bomb

        if len(bomb) == max_of_bomb:
            # 둘중에 더 적은 빨간 폭탄을 가진 것을 선택
            if count_of_red < min_of_red:
                min_of_red = count_of_red
                target = bomb
            
            if count_of_red == min_of_red:
                cy1, cx1 = get_point(bomb)
                cy2, cx2 = get_point(target)
                
                # 기준점을 이용하여 선택
                if cy1 > cy2:
                    target = bomb

                if cy1 == cy2:
                    if cx1 < cx2:
                        target = bomb
    return target


def delete_bomb(bomb):
    for y, x, in bomb:
        board[y][x] = EMPTY


def apply_gravity():

    for j in range(n):
        row_idx = n - 1

        for i in range(n - 1, -1, -1):
            if board[i][j] == STONE:
                row_idx = i - 1

            if board[i][j] != STONE and board[i][j] != EMPTY:
                while row_idx >= 0 and board[row_idx][j] == STONE:
                    row_idx -= 1

                board[row_idx][j] = board[i][j]

                if i != row_idx:
                    board[i][j] = EMPTY

                row_idx -= 1

def rotate():
    temp = [
        [EMPTY] * n
        for _ in range(n)
    ]

    for i in range(n):
        for j in range(n):
            temp[n - 1 - j][i] = board[i][j]

    return temp

answer = 0
cnt = 0
def simulate():
    global board, answer, cnt

    cnt += 1
    bombs = get_bombs()

    target_bomb = get_target(bombs)

    if len(target_bomb) <= 1:
        return False

    answer += (len(target_bomb) ** 2)

    delete_bomb(target_bomb)


    apply_gravity()


    board = rotate()

    apply_gravity()

    return True


while True:
    if not simulate():
        break

print(answer)