[APIO 2014] 수열 나누기 (Split the sequence) 풀이

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

10067번: 수열 나누기

이 문제도 몇일전에 포스트 했던 APIO 2010 특공대 문제와 비슷한 CHT (Convex Hull Technique) 응용 문제다.

 

(물론 이 문제는 DnC (Divide and Conquer) optimization 으로도 풀리나 CHT 를 활용하는 것이 더 빠르다)

 

CHT 에 관해서는 이미 APIO 2010 특공대 포스트에서 설명했으니 독자가 CHT 에 대해 익숙하다고 가정하고 설명하겠다. 만약 CHT 가 익숙하지 않다면 APIO 2010 특공대 포스트를 다시 참고하길 바란다.

 

1) $O(N^2K)$ 해법

 

아래와 같이 점화식을 정의해보자

 

$D[k][i] = $ 수열을 $k$ 개 분할하고 그 마지막 분할 위치가 $i$ 일때 최대 점수

 

그렇다면 $D[k][i]$는 아래와 같이 정의될 수 있다.

 

$$D[k][i] = max_{j}\{D[k-1][j] + (S[N]-S[i]) * (S[i] - S[j]) \}$$

 

여기서 $(S[N]-S[i]) * (S[i] - S[j])$ 는 $k - 1$번째 분할이 $j$번째 위치에서 일어났고, $k$번째 분할을 $i$번째에서 할 때 추가로 얻는 점수이다.

 

위 점화식은 $k-1$번 수열을 분할했을때의 모든 경우의수를 고려함으로써 $k$번째 분할을 $i$번째에서 할 때의 최대 점수를 구한다.

 

2) $O(NK)$ 해법

 

하지만 위 해법은 너무 느리다. 

 

이미 $N$이 최대 10만이므로 10만 제곱은 이미 2초를 아득히 넘어버릴 것이다.

 

따라서, 우리는 $N$에 대해 선형시간에 가까운 알고리즘을 고안하여야 한다.

 

하지만 위 $O(N^2K)$ 해법의 점화식의 형태를 잘 살펴보면 CHT 를 적용하기 쉬운 형태임을 알 수 있다.

 

위에서 주어진 점화식은 아래와 같이 고쳐쓸 수 있다.

 

$$D[k][i] = max_{j}\{D[k-1][j] + (S[N]-S[i]) * (S[i] - S[j]) \}$$

$$= max_{j}\{D[k-1][j] + S[N]S[i] - S[N]S[j] + S[i]^2 + S[i]S[j] \}$$

 

이는 $j$에 대한 최적화 문제이므로 $j$에 종속적이지 않은 항들은 괄호 밖으로 빼서 상수 취급 할 수 있다.

 

$$= max_{j}\{D[k-1][j] - S[N]S[j] + S[i]S[j] \} + S[N]S[i] + S[i]^2 $$

 

이제 괄호안에 있는 항 $D[k-1][j] - S[N]S[j] + S[i]S[j]$을 고려해보자.

 

$x = S[i]$로 치환하면

 

$S[j]x + (D[k-1][j] - S[N]S[j])$ 가 되며 이는 $x = S[i]$에 대한 일차함수 (직선, 선분) 이다.

 

또한 $a_i$가 음이 아닌 정수 이므로, $S[j]$는 $j$가 증가할수록 단조증가한다 (monotonically increasing).

 

따라서, 우리는 Li-Chao Tree 나 이분검색 없이 직선을 상수시간안에 관리 할 수 있다.

 

특공대 문제와 마찬가지로 덱 (deque)을 이용하여 선분의 후보들을 관리할 것이다.

 

덱에 들어 있는 선분들의 기울기가 증가하는 순으로 들어가기 때문에 현재 $x = S[i]$에 대하여 덱의 맨 앞에 들어 있는 직선 $i$와 직선 $j$ $(i < j)$의 교점이 $x$에 있거나 $x$보다 전에 있다면, 직선 $j$는 직선 $i$보다 좋고 앞으로도 더 좋을것이다. 따라서 직선 $i$를 제거한다. 이는 아래의 코드와 같이 구현 될 수 있다.

inline bool canSkip(int k, int i, int j, int threshold) {
    return (D[K(k-1)][i] - D[K(k-1)][j] + S[N] * (S[j] - S[i])) 
            <= threshold * (S[j] - S[i]);
}

// 현재 S[i] 에서 사용할 수 있는 최적의 선분을 얻는다.
while (candid.size() > 1 && canSkip(k, candid[0], candid[1], S[i])) 
	candid.pop_front();

응? 특공대 문제에서는 교점을 구하지 않았나? 싶을것이다. 사실, 

(D[K(k-1)][i] - D[K(k-1)][j] + S[N] * (S[j] - S[i])) / (S[j]-S[i])

가 교점이다. 하지만 특공대에서는 각 원소가 0보다 컸지만, 이번 문제에서는 각 원소가 0일 수도 있다. 그 말은, $S[i] = S[j]$일 수 있으며 이는 0으로 나누는 오류를 범하면서 런타임 에러가 날 것이다.

 

이를 해결하기 위해서는 나눗셈 기호를 없애주면 된다.  우리가 원하는 것은 교점이 $S[i]$보다 작거나 같은지 아닌지 하는 정보이다. 따라서, 나눗셈 항을 부등식에 양변에 곱하므로써

(D[K(k-1)][i] - D[K(k-1)][j] + S[N] * (S[j] - S[i])) <= threshold * (S[j] - S[i])

를 구할 수 있다.

 

특공대와 마찬가지로 $i$번째 선분을 추가함으로써 필요없이지는 선분 후보들을 뒤에서 제거해야한다.

 

하지만 위에서 말한 0으로 나눌 수 있는 문제 때문에 식을 이항하여 아래와 같은 함수를 정의하였다.

// intersect(k, i, j) <= intersect(k, j, l) 인지 체크.
inline bool canRemoveLastCandidate(int k, int i, int j, int l) {
    return (D[K(k-1)][i] - D[K(k-1)][j] + S[N] * (S[j] - S[i])) * (S[l] - S[j])
        <= (D[K(k-1)][j] - D[K(k-1)][l] + S[N] * (S[l] - S[j])) * (S[j] - S[i]);
}

여기서 $j$는 선분 후보에 있는 마지막 후보이고, $l$은 선분후보중 뒤에서 두번째 후보이다. 즉, $i > j > l$이다.

 

만약 $i$번째 선분과 $j$번째 선분의 교점이 $j$번째 선분과 $l$번째 번째 선분의 교점보다 전에 있다면, $j$번째 선분은 의미가 없다 ($i$번째  선분은 $j$번째 선분보다 기울기가 크므로, $i$번째 선분이 먼저 나타나면 $j$번째 선분은 절때 $i$번째 선분의 값보다 클 수 없다). 따라서, 이럴경우 제거해주면 된다.

 

이렇게 CHT 를 적용하므로써 $O(NK)$만에 문제를 해결 할 수 있다.

 

코드는 아래와 같다.

#include <algorithm>
#include <cstring>
#include <deque>
#include <iostream>
#include <vector>
#define MAXN 100000
#define MAXK 200
#define K(k) ((k)&1)
#define END (candid.size() - 1)
#define SQ(x) ((x) * (x))
typedef long long ll;
using namespace std;
// 누적합 배열, 최적해를 저장하는 배열.
// long long (8바이트) 배열을 MAXK * MAXN 개 만큼 잡으면 128MB 를 넘어간다.
// 하지만 D[k][*] 를 구할 때 D[k-1][*] 만 필요하지 그 이외는 필요하지 않으므로 두개만 잡아도 된다.
ll S[MAXN + 1], D[2][MAXN + 1];
// 역추적 배열.
int OPT[MAXK + 1][MAXN + 1];
int N, K;

// intersect(k, i, j) < threshold 인지 체크.
// 만약 이 값이 true 라면 덱 (deque) 의 첫번째 선분보다 두번째 선분이 현재 위치에서 더 좋은 후보이므로 덱의 앞에서 pop 할 수 있다.
inline bool canSkip(int k, int i, int j, int threshold) {
    return (D[K(k-1)][i] - D[K(k-1)][j] + S[N] * (S[j] - S[i])) 
        <= threshold * (S[j] - S[i]);
}

// intersect(k, i, j) <= intersect(k, j, l) 인지 체크.
inline bool canRemoveLastCandidate(int k, int i, int j, int l) {
    return (D[K(k-1)][i] - D[K(k-1)][j] + S[N] * (S[j] - S[i])) * (S[l] - S[j])
        <= (D[K(k-1)][j] - D[K(k-1)][l] + S[N] * (S[l] - S[j])) * (S[j] - S[i]);
}

// D[k][*] 를 구한다.
void dp(int k) {
    deque<int> candid;
    candid.push_back(k-1);
    // k 번째 split 이라면 자를 수 있는 처음 위치가 최소 k 이다.
    for (int i = k; i <= N; i++) {
        // 현재 S[i] 에서 사용할 수 있는 최적의 선분을 얻는다.
        while (candid.size() > 1 && canSkip(k, candid[0], candid[1], S[i])) 
        	candid.pop_front();
        int j = candid.front();
        OPT[k][i] = j;
        D[K(k)][i] = D[K(k-1)][j] + S[i] * S[j] - S[N] * S[j] + S[N] * S[i] - SQ(S[i]);
        // i 번째 선분을 집어넣으므로써 필요가 없어지는 선분들을 뒤에서 제거한다.
        while (candid.size() > 1 
        		&& canRemoveLastCandidate(k, i, candid[END], candid[END-1]))
            candid.pop_back();
        candid.push_back(i);
    }
}

// 어디서 나누면 최적의 값을 없는지 역추적한다.
vector<int> backtrace(int ind) {
    vector<int> ret;
    ret.reserve(K);
    int k = K;
    while(k > 0) {
        ret.push_back(ind);
        ind = OPT[k--][ind];
    }
    reverse(ret.begin(), ret.end());
    return ret;
}
void solve() {
    for (int k = 1; k <= K; k++) {
        memset(D[K(k)], 0, sizeof(ll) * (MAXN + 1));
        dp(k);
    }
    ll ret = 0; int ind = 1;
    for (int i = 1; i <= N; i++) if (ret < D[K(K)][i]) ret = D[K(K)][i], ind = i;

    cout << ret << '\n';
    auto opt = backtrace(ind);
    for (int x : opt) cout << x << ' ';
}
void input() {
    cin>>N>>K;
    for (int i = 1; i <= N; i++) cin>>S[i], S[i] += S[i-1];
}
int main() {
    ios::sync_with_stdio(0); cin.tie(0);
    input();
    solve();
}

여기서 한가지 주목할점은 역추적을 위해  opt 라는 int 배열을 선언한걸 볼 수 있는데 int 배열은 MAXK * MAXN 개 만큼 선언해도 약 80MB 정도로 메모리 제한안에 들어온다.