백준 21609 상어 중학교 (파이썬)

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

21609번: 상어 중학교

1. 방문을 확인하는 c를 만들되 무지개블럭이 있는 칸만 []로 저장한다

2. 현재 좌표가 일반 블럭이고 방문한적 없으면 bfs로 탐색하면서 이동가능한 칸들을 찾는다.

   bfs에는 현재 좌표와 블록 그룹의 인덱스에 해당하는 p를 입력한다.

3. bfs에서는 이동가능한 만큼 탐색한다.

   cnt 에는 블럭 그룹의 크기를 저장하고 r에는 무지개블럭 개수를 저장한다

   다음칸이 현재칸과 같은 일반 블럭이면 c에 같은 인덱스로 지정하고 cnt를 증가시킨다.

   다음칸이 무지개블럭이면 인덱스를 append하고 cnt와 r을 증가시킨다.

   탐색이 끝나면 cnt와 r을 출력한다

4. 이동 가능한 칸이 있었으면 b에 cnt, r, 좌표, p 순으로 저장하고 p를 증가시킨다

5. b에 값이 없으면 더이상 블럭 그룹이 없으므로 while문을 끝내고 답을 출력한다

6. b에 값이 있으면 정렬한다.

   정렬하면 크기가 가장 큰 블럭 그룹이 여러개 일때, 무지개블럭이 가장 많은 그룹과

   행이 가장 큰 그룹, 열이 가장 큰 그룹을 찾기 쉬워진다

7. b[-1][-1]은 규칙 1을 만족하는 그룹의 인덱스가 된다.

   일반 블럭이고 인덱스가 b[-1][-1]과 같으면 블럭을 제거하고 cnt에 개수를 센다.

   ans에 출력할 값을 갱신해준다.

8. fall 함수를 통해 블럭을 떨어뜨릴수 있는만큼 모두 떨어뜨린다.

9. list(zip(*a))[::-1] 을 사용하면 90도 반시계 회전이 된다.

   리스트 안의 값이 튜플이 되기 때문에 리스트로 바꿔준다.

10. 다시 fall 함수를 사용한다.

from collections import deque
import sys

input = sys.stdin.readline
dx = [1, -1, 0, 0]
dy = [0, 0, 1, -1]


def bfs(x, y, p):
    q.append([x, y])
    temp = a[x][y]
    c[x][y] = p
    cnt, r = 1, 0
    while q:
        x, y = q.pop()
        for i in range(4):
            nx = x + dx[i]
            ny = y + dy[i]
            if 0 <= nx < n and 0 <= ny < n:
                if a[nx][ny] == temp and c[nx][ny] == 0:
                    c[nx][ny] = p
                    cnt += 1
                    q.append([nx, ny])
                elif a[nx][ny] == 0 and p not in c[nx][ny]:
                    c[nx][ny].append(p)
                    cnt += 1
                    r += 1
                    q.append([nx, ny])
    return cnt, r


def fall(x, y):
    flag = 0
    for i in range(x+1, n):
        nx = i
        if a[i][y] > -2:
            flag = 1
            break
    if flag:
        a[nx-1][y] = a[x][y]
    else:
        a[nx][y] = a[x][y]
    a[x][y] = -2


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

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

ans = 0
while True:
    q = deque()
    c = [[0] * n for _ in range(n)]
    for i in range(n):
        for j in range(n):
            if a[i][j] == 0:
                c[i][j] = []

    p, b = 1, []
    for i in range(n):
        for j in range(n):
            if a[i][j] > 0 and c[i][j] == 0:
                cnt, r = bfs(i, j, p)
                if cnt > 1:
                    b.append([cnt, r, i, j, p])
                p += 1

    if not b:
        break

    b = sorted(b)

    cnt = 0
    for i in range(n):
        for j in range(n):
            if a[i][j] > 0 and c[i][j] == b[-1][-1]:
                a[i][j] = -2
                cnt += 1
            elif a[i][j] == 0 and b[-1][-1] in c[i][j]:
                a[i][j] = -2
                cnt += 1
    ans += cnt ** 2

    for i in range(n-2, -1, -1):
        for j in range(n):
            if a[i][j] >= 0 and a[i+1][j] == -2:
                fall(i, j)

    a = list(zip(*a))[::-1]
    a = [list(s) for s in a]

    for i in range(n-2, -1, -1):
        for j in range(n):
            if a[i][j] >= 0 and a[i+1][j] == -2:
                fall(i, j)

print(ans)