[누적합, DP] BOJ 21757 나누기

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

21757번: 나누기

 

문제 해결 알고리즘

 

수열의 누적합 배열을 만들어 준 후에

 

수열의 총합을 구하고, 총합이 0인 경우와 아닌 경우 두 가지로 나눠서 풀어줘야한다.

 

0인 경우는 마지막 0빼고 나머지 0의 개수에서 3개를 뽑는 조합으로 개수를 구해주면된다.

 

0이 아닌 경우는 다이나믹 프로그래밍을 이용해서 풀어줄 수 있는데

 

총합에서 4를 나누면 공차가 나오고 그 공차대로 늘어나는 길이 4의 수열의 개수를 구해주어야한다.

예를 들어

 

총합이 12인 수열이 있다고 하자, 그러면 4로나누면 공차는 3이면

그럼 구해줘야할 수열은 3 6 9 12이다. 이런 수열을 누적합 배열에서 찾는 과정을 DP로 풀어주면 된다.

 

자료형이 크므로 long long 자료형을 쓴다.

 

소스 코드

 

#include <bits/stdc++.h>
#define ll long long
using namespace std;

const int MAX = 100001;

ll arr[MAX], prefix_sum[MAX], sum;

int main(){
    int N; cin >> N;

    for(int i=0;i<N;i++) cin >> arr[i];

    prefix_sum[0] = arr[0];
    for(int i=1;i<N;i++) prefix_sum[i] = arr[i] + prefix_sum[i-1];
    
    sum = prefix_sum[N-1];
    ll toler = sum / 4;
    ll now = toler;

    if(toler != 0){
        ll dp[MAX];
        for(int k=0;k<4;k++){
            for(int i=0;i<N;i++){
                
                if(prefix_sum[i] == now){
                    if(now == toler) dp[i] ++;
                    else {
                        for(int j=0;j<i;j++){
                            if(now - toler == prefix_sum[j]){
                                dp[i] += dp[j]; 
                            }
                        }
                    }
                }
            
            }
            now += toler;
        }

        cout << dp[N-1];
    }
    else {
        ll dp[MAX][4];
        int cnt = 0;
        for(int i=0;i<N;i++) if(prefix_sum[i] == 0) cnt++;

        for(int i=0;i<=cnt-1;i++){
            for(int j=0;j<=3;j++){
                if(j == 0 || j == i) dp[i][j] = 1;
                else dp[i][j] = dp[i-1][j] + dp[i-1][j-1];
            }
        }
        cout << dp[cnt-1][3];
    }
    



}