[BOJ] 10067 수열 나누기

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

• 문제 요약

음이 아닌 정수 $n$개로 이루어진 수열이 있다. 이 수열을 $k+1$개로 나누어야 하는데, 한 번 나눌 때 마다 새로 나누어진 각 부분의 원소의 합을 곱한 것 만큼의 점수를 얻을 수 있다. 이 때 얻을 수 있는 점수의 최댓값과 그 순서를 구하는 문제이다.

$2 \leq n \leq 100,000$, $1 \leq k \leq \min(n-1,200)$ 이다.

• 풀이 과정

 가장 먼저 해야 할 관찰은 어느 순서로 수열을 자르던 간에 최종적으로 얻는 점수에는 변함이 없다는 것이다. 만약 어떤 부분을 2번 잘라 3개의 수열로 나누고, 이 수열의 합을 각각 $S_1$, $S_2$, $S_3$ 이라고 하자. 만약 왼쪽을 먼저 자른다고 하면, 얻을 수 있는 점수의 합은 $S_1(S_2+S_3) +S_2 S_3$이 되고, 오른쪽을 먼저 자른다고 하면 얻을 수 있는 점수의 합은 $(S_1+S_2)S_3+S_1 S_2$로 식을 전개하면 결국 $S_1 S_2 + S_2 S_3 + S_3 S_1 $로 같다는 것을 알 수 있다. 자르는 위치를 수열로 표시하면, 수열의 인접한 두 원소를 바꾸는 것에 대해 생각할 수 있다. 만약 두 원소가 서로 다른 부분을 자른다면 바꿔도 달라질 것이 없고 같은 부분을 잘라도 위에서 보인 바와 같이 결과가 바뀌지 않으므로, 어느 순서로 수열을 자르는 지는 전혀 고려하지 않아도 된다.

 따라서 다음과 같이 동적계획법 점화식을 세울 수 있다.

$DP[a][b]$=지금까지 수열의 b번째 원소까지를 a번 잘랐을 때 얻을 수 있는 최대 점수

 이 때 $DP[a][b]$는 $ 1 \leq i \le b $에 대해서 $DP[a-1][i]+S_i (S_b- S_i)$ 의 최댓값이다. 이 식을 전개하면 $S_i S_b +DP[a-1][i]-{S_i}^2 $이므로 식을 $i$에 대해 정해지는 $S_b$에 대한 일차식 꼴로 나타낼 수 있다. 이때 기울기는 $S_i$이고 y절편은 $DP[a-1][i]-{S_i}^2$인데, $S_i$는 항상 그대로이거나 증가하므로 직선의 기울기 또한 항상 증가한다고 생각할 수 있다.

 따라서 $k$ 이하의 모든 $a$에 대해서 DP를 돌리는 데, 이런 형태로 일차식 중 최대/최소를 찾아야 하는 점화식에서는 Convex Hull Trick을 사용할 수 있음이 알려져 있다. 조금 더 구체적으로 설명하면, $a$와 $b$가 고정되어 있을 때 DP값을 최대화하는 $i$를 찾고 싶다고 하자. $DP[a][i]$를 직선 형태로 나타내어 자료구조에 저장하면, DP값을 최대화 하는 $i$는 이제까지 저장된 모든 직선들에 대해 $x=S_b $에서 최댓값을 가지게 하는 $i$이다. 만약 자료구조에 직선들이 있을 때, 모든 직선 중 최댓값인 부분만을 색칠하면 색칠된 부분이 볼록 껍질 형태를 이룸을 알 수 있다. 또 원소가 추가될 때 마다 최댓값이 절대로 될 수 없는 직선들을 빼 줄 수 있음 또한 알 수 있다. 요약하자면 이와 같은 과정을 거쳤을 때 함수값은 $S_b$가 증가할 때 마다 항상 증가하거나 유지되므로 슬라이딩 윈도우와 비슷하게 최댓값을 가지는 $i$를 amortized $O(n)$으로 찾아줄 수 있다. 자세한 설명은 인터넷에 CHT를 검색해 찾아보자.

 한 가지 중요한 점은, 이 수열에는 0이 포함되어 있기 때문에 기울기가 같은 직선이 2개 연속으로 들어올 수 있다는 점이다. 일반적인 교차 x좌표를 구하는 식은 기울기가 같은 경우 (교점이 없는 경우)를 고려하지 않으므로 이 경우에는 y절편이 더 큰 직선만을 남겨야 한다. 추가로, imeimi님이 말씀하시길 이 경우에는 자르는 것이 무조건 이득이라 상관없지만 일반적인 경우에는 처음에 (0,0)이 아닌 (0,-INF)를 넣는 것이 좋다고 한다. (감사합니다!)

 그러면 남은 것은 답안을 역추적해 나가는 과정이다. 각 $DP[a][b]$에 대하여 그 DP값을 최대화하는 $i$를 저장해둘 수 있다. 이를 $pre[a][b]$라고 하면, 맨 처음에 $pre[k][n]$부터 시작하여 $pre[k][n]$의 값이 $p$일 때 $pre[k-1][p]$를 참조하고 이를 반복하는 식으로 어디를 잘라야 하는지 역추적해 나갈 수 있다. 이 문제에는 0이 있어서 같은 장소를 여러번 자르게 될 수 있으니, 만약 역추적해 나가 자른 좌표가 중복되는 경우 하나씩 오른쪽으로 밀어 주면 된다. (이 경우에, 당연하겠지만 답에 영향을 미치지 않는다. 0이니까.)

 마지막으로, 내 답안은 pre 배열을 선언하는 데 메모리를 거의 다 사용해서, dp 배열을 그대로 선언하면 MLE가 발생했다. $DP[a][b]$는 항상 $DP[a-1][i]$만 참고한다는 점에서 힌트를 얻어, dp 배열을 토글링함으로서 메모리를 절약할 수 있었다. 이와 같은 과정을 모두 구현하여 마침내 AC를 받을 수 있었다. 

• 소스 코드

#include <bits/stdc++.h>
#define MAXN 100005
using namespace std;
typedef long long ll;
typedef pair<ll,int> pli;
typedef vector<int> vi;
 
typedef struct CHT{
    ll la[MAXN], lb[MAXN]; int lz[MAXN];
    int stkn=0,pt=1;
    void insert(ll a, ll b, int z){
        while(stkn>pt){
            if(a==la[stkn] && b>=lb[stkn]) stkn--;
            else if((lb[stkn]-b)*(la[stkn]-la[stkn-1])<(a-la[stkn])*(lb[stkn-1]-lb[stkn])) stkn--;
            else break;
        }
        if(a==la[stkn] && b>=lb[stkn] && z!=0)
            return;
        la[++stkn]=a;
        lb[stkn]=b;
        lz[stkn]=z;
    }
    pli query(ll x){
        while(pt<stkn && (lb[pt]-lb[pt+1])<(la[pt+1]-la[pt])*x) pt++;
        return make_pair(la[pt]*x+lb[pt],lz[pt]);
    }
}CHT;
 
ll a[MAXN],D[2][MAXN],S[MAXN];
int pre[205][MAXN];
 
int main(){
    int n,k;
    scanf("%d %d",&n,&k);
    for(int i=1;i<=n;i++) scanf("%lld",&a[i]), S[i]=S[i-1]+a[i];
    for(int i=1;i<=k;i++){
        CHT c; c.insert(0,0,0);
        for(int j=1;j<=n;j++){
            pli temp=c.query(S[j]);
            D[i%2][j]=temp.first;
            pre[i][j]=temp.second;
            c.insert(S[j],D[(i-1)%2][j]-S[j]*S[j],j);
        }
    }
    printf("%lld\n",D[k%2][n]);
    int here=n;
    vi ans;
    ans.push_back(-1);
    for(int i=1;i<=k;i++){
        ans.push_back(pre[k-i+1][here]);
        here=pre[k-i+1][here];
    }
    sort(ans.begin(),ans.end());
    for(int i=1;i<=k;i++){
        if(ans[i]==0) ans[i]=1;
        if(ans[i]<=ans[i-1]) ans[i]=ans[i-1]+1;
        printf("%d ",ans[i]);
    }
    return 0;
}