취미로 코딩하기

※ 글쓴이는 취미로 코딩을 익혀보는 사람이라 정확하지 않은 내용을 담고 있을 수 있다 ※

이번에 볼 문제는 백준 10067번 문제인 수열 나누기이다.
문제는 아래 링크를 확인하자.

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

주어진 배열을 \(k+1\)조각냈을 때, 이 조각들이 어떤 순서로 조각났는지에 상관없이 얻는 점수는 일정함을 관찰하자.

주어진 배열(0-based index)의 \([0,n]\) 구간을 \(k\)번 나눠 얻을 수 있는 점수의 최댓값을 \(dp[k][n]\)으로 정의하자. 이 때 위의 관찰을 이용하면 \(dp\)에 대하여 다음과 같은 점화식이 성립함을 알 수 있다. (단, \(S_p\)은 주어진 배열의 \([0,p]\) 구간에 속하는 원소들의 합이다.)

\(dp[k][n] = max(dp[k-1][m] + (S_m)(S_n-S_m) (m<n)\)

위 식의 우변을 변형하면 다음과 같이 쓸 수 있다.

\(dp[k][n] = max((S_m)S_n + (dp[k-1][m] - S_{m}^{2})) (m<n)\)

\(k\)값을 고정했을 때, 모든 \(dp[k-1][n]\) 값들을 알고 있다면 위의 식에서 \(S_n\)만을 변수로 생각한다면 \(max\) 안에 들어있는 식은 \(S_m\)을 기울기로, \(dp[k-1][m] - S_{m}^{2}\)로 갖는 일차함수로 생각할 수 있고 각 일차함수들의 기울기(\(S_m\))와 대입할 변수(\(S_n\))의 값은 단조증가하므로 \(O(n)\) CHT를 이용해 모든 \(dp[k][n]\)의 값을 계산할 수 있다. 

이 모든 \(dp[k][n]\) 값들을 모두 저장하기에는 메모리가 부족하고 위의 점화식에서 각 \(k\)에 대하여 \(k-1\)에 대한 \(dp\)의 값만이 필요하므로, \(k-1\)에 대한 \(dp\)값만을 계속 저장해 문제를 해결하자.

아래는 제출한 소스코드이다.

#include <iostream>
#include <vector>
#include <deque>
using namespace std;
typedef long long ll;
typedef pair<ll, ll> pll;

ll N; int K;
deque<pair<pll, int>> dq;
ll total = 0;

ll psum[100000];
int backstep[201][100000];
ll dp[100000];

ll func(int idx, ll x) {
	return dq[idx].first.first * x + dq[idx].first.second;
}

pair<ll, int> CHT1(ll x) {
	while (dq.size() > 1) {
		if (func(0, x) <= func(1, x)) dq.pop_front();
		else break;
	}

	ll p = dq[0].first.first, q = dq[0].first.second;
	return make_pair(p * x + q, dq[0].second);
}

void CHT2(ll a, ll b, int idx) {
	while (!dq.empty()) {
		int idx = dq.size() - 1;
		ll p = dq[idx].first.first, q = dq[idx].first.second;
		if (p < a) break;
		if (b >= q) dq.pop_back();
		else return;
	}
	while (dq.size() > 1) {
		int idx1 = dq.size() - 1, idx2 = dq.size() - 2;
		ll p1 = dq[idx1].first.first, q1 = dq[idx1].first.second;
		ll p2 = dq[idx2].first.first, q2 = dq[idx2].first.second;

		if ((q2 - b) * (a - p1) < (q1 - b) * (a - p2)) break;
		dq.pop_back();
	}
	dq.push_back(make_pair(make_pair(a, b), idx));
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	cin >> N >> K;
	for (int i = 0; i < N; i++) cin >> psum[i];
	for (int i = 1; i < N; i++) psum[i] += psum[i - 1];

	for (int k = 1; k <= K; k++) {
		dq.clear();
		dq.push_back(make_pair(make_pair(0, 0), -1));
		for (int i = 0; i < N; i++) {
			pair<ll, int> ret = CHT1(psum[i]);
			CHT2(psum[i], dp[i] - psum[i] * psum[i], i);
			dp[i] = ret.first, backstep[k][i] = ret.second;
		}
	}

	N--;

	cout << dp[N] << '\n';

	vector<int> ans;

	for (int k = K; k > 0; k--) {
		ans.emplace_back(backstep[k][N] + 1);
		N = backstep[k][N];
	}

	while (!ans.empty()) {
		cout << ans.back() << ' ';
		ans.pop_back();
	}
}