[Python 3] BOJ - 21757 "나누기"

 # 문제 링크

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

21757번: 나누기

 

 # 풀이

 개인적으로 이 문제의 핵심은 누적 합과 다이나믹 프로그래밍이라고 생각한다. 우선 문제의 조건에 맞게 4등분한 구간의 합을 각각 $a_1, a_2, a_3, a_4$라고 하자. 그런데 모든 구간합의 합은 전체 수열의 합과 같다. 따라서, $a_1 + a_2 + a_3 + a_4 = 4k = s$이다. 따라서, 전체 합이 4로 나누어떨어지지 않으면 절대로 4개의 분할을 만들 수 없다.

<그림 1. 구간 분할의 경우>

 위 그림 1은 문제의 조건에 맞게 구간을 4개로 나눈 경우이다. 즉, 실선 사각형이 실제의 분할이 된다. 따라서 각 실선 사각형 내부의 합은 어떤 수 k로 같다. 그런데, 점선 사각형 내부의 합과 실선 사각형의 내부의 합을 합친다면 각각은 k, 2k, 3k, 4k이다. 그런데 점선 사각형과 실선 사각형을 합치게 되면 처음부터 시작하여 특정 시점까지의 합과 같고 이는 정확히 누적합과 같은 의미를 지닌다. 따라서, 현재 누적합이 k,2k,3k,4k라면 해당 인덱스는 분할의 경계가 될 수 있다.

<그림 2. 예제 1의 누적 합>

 위 그림 2는 예제 1의 누적합을 나타낸 그림이다. k가 3이기 때문에 3,6,9,12일 때만 분할이 가능하다. 그런데, 첫번째 그림에서는 같은 분할이라 하더라도 첫번째 6에서 긋거나 두번째 6에서 긋는 경우가 있다. 마찬가지로 두번째, 세번째 분할에서도 재귀적으로 이루어지며 subproblem들이 연결되어 있는 구조이다. 또한 첫번째 그림의 경우 첫번째 6으로 분할하는 경우와 두번째 6으로 분할하는 두 하위 경우를 통해 전체 경우를 구할 수 있으므로 다이나믹 프로그래밍의 구조를 만족한다.

 따라서 다이나믹 프로그래밍의 상태 공간을 현재 idx번째 인덱스이고 경계의 갯수가 cnt일 때, 만들 수 있는 분할의 수라고 정의하자. 만일 현재 인덱스의 누적합이 (cnt + 1) * k라면 현재 인덱스는 경계를 만들 수 있다. 그렇지 않다면 경계를 만들 수 없다. 하나 주의해야 할 점은 인덱스가 경계를 만들 수 있다고 해서 반드시 경계를 만들 필요는 없다는 것이다. 즉, 경계를 만들 수는 있으나 안 만들어도 상관없다는 것이다. 따라서 점화식은 아래와 같다. $$\begin{cases}
 & \text{ if pf[idx]  = (cnt + 1)* k,} \ \ dp[idx][cnt] += dp[idx + 1][cnt + 1] + dp[idx + 1][cnt] \\ 
 & \text{ else } \ dp[idx][cnt] += dp[idx + 1][cnt]
\end{cases}$$

 따라서, 구간이 4개라면 경계의 갯수인 cnt는 3개가 되어야 한다. 그런데, 현재 경계가 3개인데 인덱스를 벗어난다면 사실상 경계를 2개를 그은것이므로 올바른 경우가 아니기에 이것을 base case에서 예외처리를 해주어야 하는 점을 주의해야 한다

 

 # 코드

import sys

# 재귀깊이해제
sys.setrecursionlimit(200000)

# go : 현재 idx번째 인덱스이고 경계선이 cnt개일 때 유효한 분할의 갯수를 리턴하는 함수
def go(idx, cnt):
    # Base Case : 경계선이 3개인 경우
    if cnt == 3:
        # 유효한 경계선인 경우
        if idx <= n - 1:
            return 1
        # 3개를 그엇으나 실질적으로 범위밖에서 그어 유효 경계선이 2개인 경우
        return 0
    # 범위를 벗어나는 경우
    if idx == n:
        return 0
    # Memoization
    if dp[idx][cnt] != -1:
        return dp[idx][cnt]
    dp[idx][cnt] = 0
    
    # 점화식
    if pf[idx] == (cnt + 1) * k:
        dp[idx][cnt] += go(idx + 1, cnt + 1)
    dp[idx][cnt] += go(idx + 1, cnt)
    
    return dp[idx][cnt]

# 입력부
n = int(sys.stdin.readline())
arr = list(map(int, sys.stdin.readline().split()))

# 누적합 계산
pf = [0] * n
pf[0] = arr[0]
for i in range(1, n):
    pf[i] = pf[i - 1] + arr[i]
    
# 합이 4의 배수가 아닌 경우
if pf[-1] % 4 != 0:
    print(0)
    
# 합이 4의 배수인 경우
else:
    k = pf[-1] // 4
    # dp : 현재 idx번째 인덱스이고 경계선을 cnt개만큼 그엇을 때 가능한 경우의 수를 저장하는 2차원 상태공간 배열
    dp = [[-1] * 4 for _ in range(n)]
    print(go(0, 0))