BOJ-2630 색종이 만들기

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

문제 소개

최초에 한 변의 길이가 n인 정사각형 모양의 종이가 주어진다. 종이는 가로 세로 길이가 1이며 색이 하얀색 또는 파란색인 정사각형들로 이루어진다.

종이를 4등분하여 자르면 새로운 4개의 종이가 만들어지고, 이를 반복하여 만들어지는 종이의 색이 하얀색 또는 파란색이 되도록 한다.

문제 풀이

사용자는 n과 종이의 정보인 arr를 입력받고, 출력으로 누적된 각 색종이의 수인 cnt0과 cnt1를 출력한다.

최초에 입력받은 종이가 단일 색으로 이루어졌는지 확인한다. 하얀색/파란색의 색종이라면 cnt0/cnt1를 증가시키고, 색이 섞여있다면 4등분한 후 문제 페이지의 설명처럼 왼쪽위를 종이1, 오른쪽 위를 종이2, 왼쪽 아래를 종이3, 오른쪽 아래를 종이4로 하여 종이1부터 재귀적으로 단일 색으로 이루어졌는지 검사한다.

종이가 단일 색인지 확인하는 함수는 valid라는 이름으로 구현하였다. sum(arr, [])으로 list를 flatten하고 set()으로 한 종류의 원소만 가지는지 확인한다. 종이의 색이 단일이라면 색에 따라 cnt0 또는 cnt1을 증가시킨다. 단일 색이라면 True를 반환하고, 그렇지 않다면 False를 반환한다.

종이를 나누는 함수는 divide라는 이름으로 구현하였다. 리스트 컴프리헨션을 통해 나누어진 각 종이를 구하였고, 각 종이가 단일 색으로 이루어져 있는지 확인하고, 만약 색이 섞여있다면 재귀적으로 divide()를 호출한다.

코드

import sys
input = lambda: sys.stdin.readline().rstrip()


def valid(arr):
    global cnt0, cnt1
    t = set(sum(arr, []))
    if len(t) == 1:
        if t == {0}:
            cnt0 += 1
        else:
            cnt1 += 1
        return True
    else:
        return False


def divide(cur):
    size = len(cur)
    d1 = [cur[i][:size // 2] for i in range(size // 2)]
    d2 = [cur[i][size // 2:] for i in range(size // 2)]
    d3 = [cur[i][:size // 2] for i in range(size // 2, size)]
    d4 = [cur[i][size // 2:] for i in range(size // 2, size)]

    if not valid(d1):
        divide(d1)
    if not valid(d2):
        divide(d2)
    if not valid(d3):
        divide(d3)
    if not valid(d4):
        divide(d4)


n = int(input())
arr = [list(map(int, input().split())) for _ in range(n)]
cnt0 = 0
cnt1 = 0

if valid(arr):
    print(cnt0)
    print(cnt1)
    exit()

divide(arr)
print(cnt0)
print(cnt1)