BOJ 4008 - 특공대

문제 태그 : https://www.acmicpc.net/problem/4008

문제 소개

특정 값 a,b,c와 N명의 전투원들의 전투력이 주어진다. 전투원들은 인접한 전투원들끼리 묶어 몇개의 그룹으로 나눌 것이다. 각 그룹의 조정된 전투력은 그 그룹에 속한 모든 전투원의 전투력의 합을 x라고 했을 때, $ax^{2}+bx+c$이다. 이때 가질 수 있는 그룹의 조정된 전투력의 합의 최댓값을 구하시오.

문제 풀이

이 문제를 딱 보자마다 떠오르는 생각은 DP이다. '인접한 것 끼리 묶는다'도 DP를 쓰기 굉장히 좋은 조건이며, 정확히 DP[i]=1부터 i 까지 로 문제를 해결했을 때의 값으로 정의하면 깔끔하게 해결될 것 처럼 보인다. 이를 통해 DP값에 대한 식을 세워보면 다음과 같이 나온다. 
$DP[i]=\underset{j<i}{min}(DP[j]+$(j+1부터 i까지 그룹의 값)$)$ ... 1번식

이 문제를 해결하려면 $O(N^{2})$의 시간복잡도가 필요하다. 따라서 DP[j]+(j+1부터 i까지 그룹의 값)에 좀 더 집중을 해서 관찰을 해야한다. 일단 (j+1부터 i까지 그룹의 값)을 구하기 위해서는 j+1부터 i까지의 부분 합이 필요하다. 하지만 전처리를 해도 구간이 $N^2$ 스케일의 갯수가 나오게 된다.
이 문제를 해결하면서 j+1과 j의 연관성을 줄 수 있는 방법은 prefix sum을 이용하는 것이다. 1부터 i까지의 prefix sum을 $pre[i]$라고 하면 j+1부터 i까지의 부분 합은 $pre[i]-pre[j]$가 될 것이다. 이를 이용하여 1번식을 다시쓰면,

$DP[i]=\underset{j<i}{min}(DP[j]+a(pre[i]-pre[j])^{2}+b(pre[i]-pre[j])+c)$라는 식이 나오게 된다.
이렇게 하면 구간합은 해결이 되지만, 임의의 i,j 쌍에 대해 모두 확인해야 하므로 여전히 시간복잡도는 $O(N^2)$을 가지게 된다. 하지만 식을 좀 더 정리해 보면 CHT를 쓸 수 있는 형태라는 것을 알 수 있다.

$(-2\cdot a\cdot pre[j]+ b)\cdot pre[i] +a\cdot pre[j]^2 +c+a\cdot pre[i]^2 $
위와 같은 식이 나오게 되고, pre[i]의 배수이면서 j와 관련된 부분, pre[i]와 관련이 없으면서 j와 관련된 부분, 그 밖의 부분으로 분리가 되고, 이는 $ax+b+c$ 꼴을 갖추게 된다. c는 추가적으로 i와 관련해서 더해주는 값이기 때문에 CHT를 이용하는데 지장이 없다. 

따라서 CHT를 이용하여 $O(NlgN)$에 이 문제를 해결 할 수 있다. 

코드

#include <bits/stdc++.h>
using namespace std;
typedef long long int ll;
 
struct lin{  
    ll a,b;
};
 
ll sz=0;
lin stk[1010101];
double cx(ll a,ll b) {
    return 1.0 * (double)(stk[a].b - stk[b].b) / (stk[b].a - stk[a].a);
}
void insert(lin v){
    stk[sz] = v;
    while(1<sz&&cx(sz-2,sz-1)>cx(sz-1,sz)){
        stk[sz-1] = stk[sz];
        sz--;
    }
    sz++;
}
ll sol(ll x){ 
    ll lo=0,hi=sz-1;
    while(lo<hi){
        ll mid=(lo+hi)/2;
        if (cx(mid,mid+1)<=x) lo=mid+1;
        else hi=mid;
    }
    return stk[lo].a*x+stk[lo].b;
}
void clear(){
    sz=0;
}
 
ll x[1010101];
ll pre[1010101];
ll dp[1010101];
 
int main(){
    ll a,b,c,n,i,j,k;
    scanf("%lld",&n);
    scanf("%lld %lld %lld",&a,&b,&c);
    for(i=1;i<=n;i++){
        scanf("%lld",&x[i]);
        pre[i] = pre[i-1]+x[i];
    }
    lin l;
    l.a = b;
    l.b = 0;
    insert(l);
    for(i=1;i<=n;i++){
        dp[i] = c+a*pre[i]*pre[i]+sol(pre[i]);
        l.a = -2*a*pre[i]+b;
        l.b = dp[i]-b*pre[i]+a*pre[i]*pre[i];
        insert(l);
    }
    printf("%lld",dp[n]);
    return 0;
}