[BOJ/python]21757번 나누기

백준 21757번 나누기 파이썬

문제 링크 - https://www.acmicpc.net/problem/21757

21757번: 나누기

누적합과 dp를 활용하여 풀어야 했는데 dp 점화식을 세우는 것이 생각보다 쉽게 떠오르지 않아서 당황했다. 현재 풀이보다 dp부분을 좀 더 깔끔한 코드로 작성할 수 있을 것 같다.

처음 시도할 때는 구간합 배열에서 크기가 같은 구간이 되도록 나누어주는 divide 함수를 구현하여 풀고자 했다.

즉, 구간합 배열을 2등분하고 다시 각각의 쪼개진 구간을 반으로 나누어 총 개수를 세고자 시도했지만 당연하게도 시간초과가 발생했다.

 

0. 오답 코드

N = int(input())
A = list(map(int, input().split()))

def divide(prefix, before):
    idx = []
    for i in range(len(prefix)-1):
        if prefix[-1] - prefix[i] == prefix[i] - before:
            idx.append(i)
            
    return idx

def getPrefix(A): #O(N)
    prefix = [0 for _ in range(len(A))]
    prefix[0]=A[0]
    for i in range(1,len(A)):
        prefix[i] = prefix[i-1] + A[i]    

    return prefix

prefix1 = getPrefix(A)
idx1 = divide(prefix1, 0 )

sol=0
for i in idx1:  #O(N)
    if i>0 and i<len(A)-1:
        tmp1 = divide(prefix1[:i+1],0)
        tmp2 = divide(prefix1[i+1:], prefix1[i])
        sol += len(tmp1)*len(tmp2)
print(sol)

1. 풀이

처음부터 구간합을 4등분하여 풀어야한다. 이 때 전체 배열의 합이 4의 배수인 경우와 4의 배수가 아닌 경우가 있으므로 각각 경우를 나누어 생각한다.

우선 주어진 수열의 구간합 배열을 계산한다.

1. 전체 배열의 합이 4의 배수가 아닌 경우
   • 임의의 구간 a,b,c,d가 있을 때 a+b+c+d가 prefix[-1]에 저장된다. 만약 임의의 구간 a,b,c,d의 합이 같다면 prefix[-1]의 값은 sum(a)*4 이므로 4의 배수인 경우에만 문제를 해결할 수 있다.
   • 즉 불가능한 경우이므로 0을 출력한다.
2. 전체 배열의 합이 0인 경우
   • a, b, c, 0, d, 0, e, 0, f, 0 와 같은 구간합 배열을 생각해보자. 가능한 경우는 0을 기준으로 4등분 된 1가지 밖에 없다.
   • a, b, 0, c, 0, d, 0, e, 0, f, 0 에서는 총 4가지 경우가 가능하다.
   • 즉 0의 개수가 tmp개라면 조합을 이용하여 C(tmp-1,3)라는 답을 쉽게 구할 수 있다.
3. 전체 배열의 합이 4의 배수인 경우
   • 전체 배열의 합이 12라고 하자. 이 때 구간합을 4등분하면 모든 구간합의 크기는 3일 것이다. 즉 구간합 배열 자체에서는 첫번째 구간의 값은 3, 두번째 값은 6, 세번째 값은 9, 마지막값은 12이다.
   • 즉 구간합 배열에서 3,6,9,12가 순서대로 등장하는 경우가 4등분이 가능한 경우이다.
   순서대로 등장하는 조합을 모두 세기 위해 DP를 활용해서 풀면 된다.
   • 구간합 배열을 처음부터 끝까지 탐색하며 해당 인덱스’까지’ ‘가능한’ 3,6,9,12의 개수를 갱신하기 위해 2차원 배열을 사용한다. dp[i]=[0,0,0,0]
   • 3은 가장 첫번째 수이므로 3이 등장하면 이전 인덱스의 3의 개수에 +1을 해준다. dp[i][0] = dp[i-1][0]+1
   • 6이 등장하면 이전 인덱스까지 등장한 3에 6을 그대로 붙이면 가능한 순서이다. 여기에 이전 인덱스까지 가능했던 3-6 순서의 조합의 개수를 더해주면 된다. dp[i][1] = dp[i-1][0] + dp[i-1][1]
   • 9가 등장했을 때는 마찬가지로 dp[i][2] = dp[i-1][1] + dp[i-1][2]
   • 구간합 배열에서 마지막 인덱스가 아닌 중간 12가 등장하면 4등분 된 경우가 아니다. 즉 12가 중간에 등장한 경우는 계산할 필요가 없다.
   • 모든 구간합 배열을 탐색한 후 마지막 인덱스에서 9의 개수를 통해 가능한 경우의 수(3-6-9-12)를 모두 계산할 수 있다.

2. 코드

N = int(input())
A = list(map(int, input().split()))

def getPrefix(A): #O(N)
    prefix = [0 for _ in range(len(A))]
    prefix[0]=A[0]
    for i in range(1,len(A)):
        prefix[i] = prefix[i-1] + A[i]    
    return prefix

prefix1 = getPrefix(A)

if prefix1[-1] % 4 != 0:
    print(0)

elif prefix1[-1]==0:
    tmp=0
    for i in range(N):
        if prefix1[i]==0:
            tmp+=1
    
    sol = (tmp-1)*(tmp-2)*(tmp-3)//6
    print(sol)
    
else:   
    div = prefix1[-1]//4
    
    dp = [ [0 for _ in range(4)] for _ in range(N) ]
    
    for i in range(1, N):
        n = prefix1[i]
        if n == div:
            dp[i][1] = dp[i-1][1]
            dp[i][2] = dp[i-1][2]
            dp[i][3] = dp[i-2][3]
            
            dp[i][0]=dp[i-1][0]+1
            
            
        elif n == div*2:
            dp[i][0] = dp[i-1][0]
            dp[i][2] = dp[i-1][2]
            dp[i][3] = dp[i-2][3]
            
            dp[i][1] = dp[i-1][1]+dp[i-1][0]
            
        elif n == div*3:
        
            dp[i][0] = dp[i-1][0]
            dp[i][1] = dp[i-1][1]
            dp[i][3] = dp[i-2][3]
            
            dp[i][2] = dp[i-1][2]+dp[i-1][1]
            
            
        # elif n == div*4:  # 이것 때문에 틀렸다!
        #     dp[i][0] = dp[i-1][0]
        #     dp[i][1] = dp[i-1][1]
        #     dp[i][2] = dp[i-2][2]
            
        #     dp[i][3] = dp[i-1][3]+dp[i-1][2]
            
            
        else:
            dp[i]= dp[i-1]  
    print(dp[N-1][2])