본문 바로가기

개발/algorithm

[백준 21609번] 상어 중학교 - python

문제 링크

 

풀이

나름 간단해보이는데, 난이도가 높다면 함정이 있다는 것 .. 

여기서 주의해야할 점은 

일반 블록 색은 모두 같아야하며, 무지개 블록은 상관없다는 것이다. 

따라서 가장 큰 블록을 bfs를 사용해서 구할 때, visited 배열에 방문표시를 하게 되는데, 계산이 끝나면 무지개 블록은 다시 방문 표시를 해제해 주어야 한다... 

 

사용한 주요 알고리즘은 다음과 같다. 

1. 시계 방향 90도 회전하는 함수 

def rotate():
    global graph
    new = [ [0] * n for _ in range(n)]
    
    for i in range(n):
        for j in range(n):
            new[n-j-1][i] = graph[i][j]
    
    graph = new

시계 방향 90도 회전은 인덱스 변화가

new_x = origin_y

new_y = n-1-origin_x 이고,

반시계 방향 90도 회전은 인덱스 변화가

new_x = n-1-origin_y

new_y = origin_x 임을 기억하자 

 

2. 중력 함수 구현 

각 세로 줄마다, 아래 위치에서 부터 위로 올라가며 다음 과정을 반복해준다. 

1. 현재 위치에서 아래 블록을 확인한다. 

2. 아래 블록이 -1 거나 다른 블록이 있다면 중단한다.

- 여기서 다른 블록이 있을 수 있는 경우는 제일 아래에 블록이 있는 경우이다. 

3. 빈칸이라 내려갈 수 있다면 현재 블록을 아래로 내리고, 현재 위치를 빈칸으로 바꿔준다. 

4. 중단될 때까지 아래로 내려가며 위 과정을 반복한다. 

def gravity():
    
    for j in range(n):
        for i in range(n-2, -1, -1):
            if graph[i][j] >= 0 :
                r = i+1
                while True :
                    if r >= n :
                        break 
                    if graph[r][j] == -1 or graph[r][j] >= 0 :
                        break 
                
                    graph[r][j] = graph[r-1][j]
                    graph[r-1][j] = -2 
                    r += 1

 

전체적인 알고리즘은 다음과 같다. 

1. bfs로 가장 큰 블록찾기 

- heap에 (-블록 크기, -무지개 블록 개수, -행 번호, - 열번호)를 넣어주어 조건에 맞는 가장 큰 블록을 구한다.

2. 찾은 블록이 없다면 중단

3. 가장 큰 블록을 지운다. 

- 이 때 빈칸은 -2로 표시한다.

4. 중력 함수를 실행한다.

5. 반 시계 방향 90도 회전 함수를 실행한다.

6. 중력 함수를 실행한다.

# https://www.acmicpc.net/problem/21609
from collections import deque 
import heapq 
dx = [1, -1, 0, 0]
dy = [0,0,1,-1]
ans = 0

# 블록 크기 찾는 함수 
def bfs(a, b):
    q = deque()
    q.append((a,b))
    
    heap = []
    
    # 기준 블록을 찾기 위해 포함된 블록들의 x,y 좌표를 힙에 대입한다. 
    heapq.heappush(heap, (a,b))
    # 블록 개수 
    cnt = 1 
    # 무지개 블록 개수 
    rainbow = 0
    # 방문 표시 
    visited[a][b] = True 
    # 같은 색깔 
    color = graph[a][b]
    
    # 무지개 블록 위치 
    rainbow_idx = []
    while q :
        x, y = q.popleft()

        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            
            if nx < 0 or nx >= n or ny < 0 or ny >= n :
                continue 
            
            # 같은 색깔 블록이라면 
            if graph[nx][ny] == color and not visited[nx][ny]:
                cnt += 1 
                visited[nx][ny] = True 
                heapq.heappush(heap, (nx,ny))
                q.append((nx,ny))
                
            # 무지개 블록이라면 
            if graph[nx][ny] == 0 and not visited[nx][ny]:
                cnt += 1 
                visited[nx][ny] = True 
                rainbow += 1 
                rainbow_idx.append((nx,ny))
                q.append((nx,ny))
                
    # 무지개 블록 방문 표시 해제 
    for a,b in rainbow_idx:
        visited[a][b] = False 
        
    # 기준 블록 행열 번호 
    stand_x, stand_y = heapq.heappop(heap)

	# 블록 개수 -> 무지개 블록 개수 -> 행 번호 -> 열 번호 순 
    # 큰 값 우선이므로 마이너스를 붙여준다. (최대힙)
    result = [-cnt, -rainbow, -stand_x, -stand_y]
    
    return result
    
# 블록 제거 함수 
def remove(x, y):
    global ans 
    q = deque()
    q.append((x,y))
    
    visited = [[False] * n for _ in range(n)]
    visited[x][y] = True 

    color = graph[x][y]
    check_cnt = 1 
    graph[x][y] = -2 
    while q :
        x ,y = q.popleft()
        
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            
            if nx < 0 or nx >= n or ny < 0 or ny >= n :
                continue 
            
            if (graph[nx][ny] == color or graph[nx][ny]==0)and not visited[nx][ny]:
                graph[nx][ny] = -2 
                visited[nx][ny] = True 
                check_cnt += 1 
                q.append((nx,ny))
                
    # 제거된 블록 개수 제곱 더하기             
    ans += check_cnt ** 2 
    
# 중력 함수 
def gravity():
    
    for j in range(n):
        for i in range(n-2, -1, -1):
            if graph[i][j] >= 0 :
                r = i+1
                while True :
                	# 범위 벗어나면 중단 
                    if r >= n :
                        break 
                    # 검정 블록이나 다른 블록을 만나면 중단 
                    if graph[r][j] == -1 or graph[r][j] >= 0 :
                        break 
                	# 아래로 내려준다 
                    graph[r][j] = graph[r-1][j]
                    graph[r-1][j] = -2 
                    r += 1 
 # 반시계 방향 90도 회전 함수
def rotate():
    global graph
    new = [ [0] * n for _ in range(n)]
    
    for i in range(n):
        for j in range(n):
            new[n-j-1][i] = graph[i][j]
    
    graph = new 
                    
                
                         
n, m = map(int,input().split())

graph = []

for _ in range(n):
    graph.append(list(map(int,input().split())))
    
while True :
    visited = [[False] * n for _ in range(n)]
    
    # 블록 저장 최대 힙 
    h = []
    for i in range(n):
        for j in range(n):
            if graph[i][j] > 0 and not visited[i][j]:
                tmp = bfs(i,j)
                # 블록 크기가 최소 2이상 
                if -tmp[0] < 2 :
                    continue 
                heapq.heappush(h, tmp)              
    # 블록 없으면 중단 
    if len(h) == 0:
        break 
    
    # 가장 큰 블록의 기준 블록 위치 찾기 
    tmp = heapq.heappop(h)
    s_x = -tmp[2]
    s_y = -tmp[3]
    
    # 블록 제거 
    remove(s_x, s_y) 
      
    gravity()  
    rotate()
    gravity()

print(ans)