백준 11559. Puyo Puyo

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

import sys
from collections import deque

input = sys.stdin.readline


arr = [list(input().strip()) for _ in range(12)]
ans = 0  # 연쇄 횟수


# BFS로 4개 이상의 같은 색 뿌요를 찾는 함수
def bfs(si, sj, visited):
    q = deque([(si, sj)])
    visited[si][sj] = True
    puyo_list = [(si, sj)]
    color = arr[si][sj]  # 뿌요 색상

    while q:
        i, j = q.popleft()
        for di, dj in [(1, 0), (-1, 0), (0, 1), (0, -1)]:  # 상하좌우 탐색
            ni, nj = i + di, j + dj
            if 0 <= ni < 12 and 0 <= nj < 6 and not visited[ni][nj] and arr[ni][nj] == color:
                visited[ni][nj] = True
                q.append((ni, nj))
                puyo_list.append((ni, nj))

    # 4개 이상 연결된 경우 반환
    return puyo_list if len(puyo_list) >= 4 else []


# 중력 적용 함수
def apply_gravity():
    for y in range(6):
        for x in range(10, -1, -1): # 맨아래 한칸 남겨둬야하니까 10~
            for k in range(11, x, -1): # 본인의 아래부터 맨아래까지 빈칸 있는지 확인. 맨아래까지 탐색하니까 11~
                if arr[x][y] != "." and arr[k][y] == ".": # 나는 빈칸이 아닌데 내 한칸밑이 빈칸이다?
                    # 원래 내자리는 빈칸으로 만들고 나는 한칸 내려가기
                    arr[k][y] = arr[x][y]
                    arr[x][y] = "."

while True:
    visited = [[False] * 6 for _ in range(12)]  # 방문 체크
    is_burst = False
    # 터뜨릴 블록 찾기
    for i in range(12):
        for j in range(6):
            if arr[i][j] != '.' and not visited[i][j]:
                puyo_list = bfs(i, j, visited)
                if puyo_list:  # 터뜨릴 블록이 있다면 추가
                    for x, y in puyo_list:
                        is_burst = True
                        arr[x][y] = "."

    if is_burst:
        apply_gravity()
        ans += 1
    else:
        break

print(ans)

내 처음 코드

import sys
from collections import deque
input = sys.stdin.readline

arr = [list(input().rstrip()) for _ in range(12)]
ans = 0

def main():
    global ans, arr
    is_burst = 0

    def bfs(si,sj):
        visited = [[0]*6 for _ in range(12)]
        q = deque([(si,sj)])
        visited[sj][si] = 1
        cnt = 1
        lst = [(si,sj)]
        while q:
            i, j = q.popleft()
            for di, dj in [(1, 0), (-1, 0), (0, 1), (0, -1)]:
                ni, nj = i + di, j + dj
                if 0 <= ni < 6 and 0 <= nj < 12 and not visited[nj][ni] and arr[nj][ni] == arr[j][i]:
                    visited[nj][ni] = visited[j][i] + 1
                    q.append((ni, nj))
                    lst.append((ni, nj))
                    cnt += 1
        if cnt >= 4:
            return lst

    burst = []
    def find_burst():
        for i in range(6):
            for j in range(12):
                if arr[j][i] != '.':
                    burst_lst = bfs(i,j)
                    if burst_lst:
                        burst_lst.sort(key=lambda x : x[1])
                        burst.append(burst_lst)
                        return 1
    is_burst = find_burst()

    for bur in burst:
        for i, j in bur:
            for k in range(12):
                if arr[j-k-1][i] != '.':
                    arr[j-k][i] = arr[j-k-1][i]

                else:
                    arr[j-k][i] = '.'
                    break

    if is_burst:
        ans += 1
    else:
        return -1

while True:
    s = main()
    if s == -1:
        break

print(ans)

터지는 것과 동시에 이동을 처리하려고 하니까 굉장히 복잡해지고 난리났다...

일단 '.'으로 다 터뜨린 다음 중력 적용을 했어야 됐는데...