[Python/파이썬] 백준 21757번 나누기

21757번: 나누기

 

문제

 $N$개의 정수 수열 $A_1, A_2, \dots , A_N$이 주어진다. 수열을 각각이 연속된 네 부분으로 나누려고 한다. 단, 각 부분은 최소 하나의 수를 포함해야 한다. 또, 각 부분의 합은 모두 같아야 한다. 즉, 어떤 $i, j, k$ ($1 \le i < j < k < N$)에 대해서$[A_1, \dots, A_i], [A_{i+1}, \dots, A_j], [A_{j+1}, \dots, A_k], [A_{k+1}, \dots, A_N]$으로 나눈다.

예를 들어 주어진 수열이 $4, −1, 2, 1, −3, 1, 2, 2, 1, 3$이라고 하자. 이 수열을 아래와 같이 나누면 각 부분의 합이 달라서 허용되는 형태가 아니다.

$[4, −1, 2], [1, −3, 1, 2], [2, 1], [3]$ 

아래과 같이 나눈 경우 각 부분의 합이 모두 같다.

$[4, −1], [2, 1], [−3, 1, 2, 2, 1], [3]$ 

아래와 같이 나눈 경우들도 각 부분의 합이 모두 같다.

$[4, −1], [2, 1, −3, 1, 2], [2, 1], [3]$ 혹은 $[4, −1, 2, 1, −3], [1, 2], [2, 1], [3]$ 

수열을 입력 받아 위와 같이 나눌 수 있는 가능한 방법의 개수를 계산하는 프로그램을 작성하라.

입력

첫 번째 줄에 수열의 길이 $N$이 주어진다.

두 번째 줄에 $N$개의 정수 $A_1, A_2, \dots , A_N$이 공백 하나씩을 사이로 두고 주어진다.

출력

첫 번째 줄에 가능한 방법의 개수를 출력한다.

출력 값이 매우 클 수 있으므로 C, C++ 언어에서는 long long 형의 변수를, Java에서는 long 형의 변수를 사용해야 한다.

코드

import sys
input = sys.stdin.readline

n = int(input())
arr = list(map(int, input().split()))

# 누적합
preSum = [0] * (n+1)
for i in range(1, n+1):
    preSum[i] = preSum[i-1] + arr[i-1]

ans = 0

if preSum[n] % 4 == 0:
    partial_sum = preSum[n] // 4
    dp = [[0]*4 for _ in range(n+1)]
    dp[0][0] = 1
    for i in range(1, n+1):
        dp[i][0] = 1
        for j in range(1, 4):
            dp[i][j] = dp[i-1][j]
            if preSum[i] == partial_sum * j:    # 이번에 자르기 가능한 경우
                dp[i][j] += dp[i-1][j-1]
    ans = dp[n-1][3]
print(ans)

수들의 총합이 4로 나눠 떨어지는 경우에만 수열을 각 부분의 합이 같게 네 부분으로 나눌 수 있다.

 

네 부분으로 나눌 수 있는 경우에는 다이나믹 프로그래밍을 이용하여 i번째 수까지 j번 자를 수 있는 방법의 수를 dp 배열에 저장할 것이다. partial_sum은 네 부분수열으로 잘랐을 때 각 수열이 가져야 하는 총합이다.

 

이제 i부터 n까지 반복문을 돌며 dp 배열의 값을 채워준다. preSum[i]가 partial_sum*j와 같은 경우에는 이번에 배열을 자를 수 있다는 뜻이다. 그렇기 때문에 i-1번째 수까지 j-1번 자를 수 있는 방법의 수인 dp[i-1][j-1]을 dp[i][j]에 더해준다.

 

 

처음에는 누적합+DFS로 풀었다가 7점 맞고...DP를 써야할 것 같다는 생각이 들어서 풀이 방향을 바꿨다. 이 문제에서 포인트는 누적합을 사용해야 한다는 점인 것 같다.