코드트리 예술성 (파이썬)

https://www.codetree.ai/training-field/frequent-problems/problems/artistry

코드트리 | 코딩테스트 준비를 위한 알고리즘 정석

import sys
from collections import deque

input = sys.stdin.readline

# 인접한 변의 수를 구할 때를 고려해서 우하 방향을 앞에둠
dx = [0, 1, 0, -1]
dy = [1, 0, -1, 0]

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

ans = 0
for r in range(4):
    # bfs로 그룹을 만들고 group에 각 그룹의 번호를 매긴다
    # group_cnt: 그룹에 속한 칸의 수
    # group_num: 그룹을 이루고 있는 숫자 값
    # cnt: group_cnt에 넣을 값 = 그룹 번호
    # num: group_num에 넣을 값 = 해당 그룹의 값
    group = [[0 for _ in range(n)] for _ in range(n)]
    group_cnt, group_num = dict(), dict()
    q = deque()
    num = 1
    for i in range(n):
        for j in range(n):
            if group[i][j] == 0:
                q.append([i, j])
                group[i][j] = num
                group_num[num] = a[i][j]
                cnt = 1
                while q:
                    x, y = q.popleft()
                    for k in range(4):
                        nx = x + dx[k]
                        ny = y + dy[k]
                        if 0 <= nx < n and 0 <= ny < n:
                            if group[nx][ny] == 0 and a[nx][ny] == a[x][y]:
                                q.append([nx, ny])
                                group[nx][ny] = num
                                cnt += 1
                group_cnt[num] = cnt
                num += 1

    # 서로 다른 두 그룹이 맞닿은 변 개수 구하기
    # connect에 2차원 배열 형태로 저장
    # ex) connect[1][2]는 그룹1과 그룹2가 맞닿은 변 개수 값
    # 우하 방향만 탐색하면 중복을 고려하지 않아도 됨
    connect = [[0 for _ in range(num)] for _ in range(num)]
    for x in range(n):
        for y in range(n):
            for i in range(2):
                nx = x + dx[i]
                ny = y + dy[i]
                if 0 <= nx < n and 0 <= ny < n:
                    if group[nx][ny] != group[x][y]:
                        g1 = group[x][y]
                        g2 = group[nx][ny]
                        connect[g1][g2] += 1
                        connect[g2][g1] += 1

    # 모든 조화로움 점수를 ans에 더해준다
    for i in range(1, num):
        for j in range(i, num):
            if connect[i][j] != 0:
                a_cnt, b_cnt = group_cnt[i], group_cnt[j]
                a_num, b_num = group_num[i], group_num[j]
                ans += (a_cnt + b_cnt) * a_num * b_num * connect[i][j]

    # 3회전을 마쳤으므로 정답을 출력하고 break
    if r == 3:
        print(ans)
        break

    # 가운데 십자가 형태의 배열을 반시계 회전
    temp_12 = [a[i][n // 2] for i in range(n // 2)]
    for i in range(n//2):
        a[i][n // 2] = a[n // 2][n - 1 - i]
        a[n // 2][n - 1 - i] = a[n - 1 - i][n // 2]
        a[n - 1 - i][n // 2] = a[n // 2][i]
        a[n // 2][i] = temp_12[i]

    # 1사분면의 정사각 배열 시계 회전
    rotate = [a[i][n // 2 + 1:] for i in range(n // 2)]
    rotate = list(zip(*rotate[::-1]))
    for i in range(n // 2):
        a[i][n // 2 + 1:] = rotate[i]

    # 2사분면의 정사각 배열 시계 회전
    rotate = [a[i][:n // 2] for i in range(n // 2)]
    rotate = list(zip(*rotate[::-1]))
    for i in range(n // 2):
        a[i][:n // 2] = rotate[i]

    # 3사분면의 정사각 배열 시계 회전
    rotate = [a[i][:n // 2] for i in range(n // 2 + 1, n)]
    rotate = list(zip(*rotate[::-1]))
    for i in range(n // 2 + 1, n):
        a[i][:n // 2] = rotate[i - (n // 2 + 1)]

    # 4사분면의 정사각 배열 시계 회전
    rotate = [a[i][n // 2 + 1:] for i in range(n // 2 + 1, n)]
    rotate = list(zip(*rotate[::-1]))
    for i in range(n // 2 + 1, n):
        a[i][n // 2 + 1:] = rotate[i - (n // 2 + 1)]