[BOJ] 백준 22355. 말뚝 (Platinum II)

UCPC 2021 예선 E번으로 나왔던 문제이다.

대회 당시에 lazy segment tree 느낌이 팍팍 들어서 일단 보류하고 다른 문제로 넘어갔던 기억이 나는데, 넘어가길 잘했던 문제. 상당히 시간이 오래 걸리며, 구현량도 꽤 있는 편이었다.

또한, 그래프가 unimodal하고, 기울기가 같은 구간이 존재하므로 ternary_search (삼분탐색)을 이용하기까지 해야 하는 문제였다. 이분탐색으로도 가능하다는데, 잘은 모르겠다.

 

문제는 아래와 같다.

 

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

22355번: 말뚝

---

의식의 흐름 및 해설

먼저, 아름다움이 K 이상이라고 말하지만, 아름다움이 K일 때 말뚝을 건드려야할 것이 가장 적기 때문에 K일 때가 답임을 알 수 있다.

 

따라서 아름다움이 K일 때 구간 [1, K], [2, K+1], ... ,[N-K+1, N]일 때를 스위핑으로 살펴보고 가장 적은 힘이 드는 경우를 출력해주면 된다.

 

구간 [1, K]일 때 높이를 H[1], H[2], ... , H[K] 중 어떤 높이로 평평하게 만들어야 가장 아름다울지 일일이 비교하면 시간복잡도는 O(NK) 또는 그 이상으로 TLE(Time Limit Exceeded, 시간초과)를 받게 될 것이다. 조금 생각을 더 해보자. 말뚝을 올릴 때와 내릴 때 드는 힘이 일정하므로 각 높이에 따라 말뚝을 움직이는 데에 필요한 힘을 일차함수 꼴로 만들 수 있다. 

 

예제 입력2에 해당되는 경우를 예시로 들어보겠다.

5 3
1 2 3 2 1
1 3 1 3 4
1 3 5 3 1

예를 들어 세 번째 말뚝은 현재 높이가 3이며, 말뚝을 박는 데에 5만큼, 말뚝을 들어올리는 데에 1만큼 힘이 소요된다.

따라서 세 번째 말뚝에 소요되는 힘은 아래와 같이 나타낼 수 있다.

 말뚝의 높이를 x로 만드는 데 소요되는 힘

 구간 [1, K]에 해당되는 각 높이에 따른 말뚝에 소요되는 힘은 아래와 같다. 

구간 [1, K]의 모든 말뚝에 소요되는 힘을 합쳐야 하므로 모두 더해준다.

이 과정을 lazy 구간합 segment tree로 처리해주어야 로그시간복잡도로 처리할 수 있다.

이 부분은 코드로 아래와 같다.

update(tree, 1, 1, MAX, 1, a[i]-1, -desc[i]);
update(cons, 1, 1, MAX, 1, a[i]-1, desc[i] * a[i]);
update(tree, 1, 1, MAX, a[i], MAX, asc[i]);
update(cons, 1, 1, MAX, a[i], MAX, -asc[i] * a[i]);

일차함수 꼴(y=ax+b)이기 때문에 기울기의 합을 나타내는 세그먼트트리인 tree, 상수의 합을 나타내는 세그먼트트리인 cons 이렇게 두 개의 세그먼트트리로 관리하였다. 높이를 1 미만으로 평평하게 하는 경우는 최소의 힘이 아닐 것이 자명하고, 일차함수 꼴이 a[i]를 기점으로 V자 꼴로 절댓값 그래프 모양을 지니므로 [1, a[i]), [a[i], 10만] 구간으로 나눠서 update해야 함에 주의하자. Range Update이므로 레이지 세그먼트트리를 활용했다.

 

 

위 함수식을 그래프로 나타내면 아래와 같다.

따라서 구간 [1, K] 의 말뚝을 평평하게 만드는데 드는 최소의 힘은 높이가 3일 때 5의 힘이 들 때가 최소임을 알 수 있다. 그래프를 보면 평평한 구간이 존재하기 때문에 Unimodal하긴 하지만, 삼분탐색으로 최솟값을 찾아야 함을 알 수 있다. 그 코드는 아래와 같다.

ll solve(int s, int e) {
	while (s + 3 <= e) {
		int p = (s * 2 + e) / 3;
		int q = (s + e * 2) / 3;
		ll A1 = query(tree, 1, 1, MAX, p, p);
		ll A2 = query(tree, 1, 1, MAX, q, q);
		ll B1 = query(cons, 1, 1, MAX, p, p);
		ll B2 = query(cons, 1, 1, MAX, q, q);
		if (A1 * p + B1 > 
			A2 * q + B2)
			s = p;
		else
			e = q;
	}
	ll ret = LNF;
	for (int i = s; i <= e; i++) {
		ll A = query(tree, 1, 1, MAX, i, i);
		ll B = query(cons, 1, 1, MAX, i, i);
		ret = min(ret, A * i + B);
	}
	return ret;
}

 

나머지 구간 [2, K+1], ... , [N-K+1, N] 에서도 동일하게 수행하여 그 중 최솟값을 골라내면 된다.

---

코드

#include <iostream>
#include <algorithm>
#include <string>
#include <vector>
#include <cstring>
#include <queue>
#include <cmath>
#include <stack>
#include <set>
#include <map>
#define all(v) (v).begin(), (v).end()
#define press(v) (v).erase(unique(all(v)), (v).end())
using namespace std;
typedef long long ll;
typedef pair<int, int> pi;
typedef pair<int, pi> pii;
const int MAX = 100011;
const int INF = 0x3f3f3f3f;
const ll LNF = 1e18;
struct Tree {
	ll val, lazy;
};
ll N, K, a[MAX], asc[MAX], desc[MAX], ans = LNF;
vector<Tree> tree, cons;
void update_lazy(vector<Tree>&tree, int node, int s, int e) {
	if (tree[node].lazy) {
		tree[node].val += ((ll)e - s + 1) * tree[node].lazy;
		if (s != e) {
			tree[node * 2].lazy += tree[node].lazy;
			tree[node * 2 + 1].lazy += tree[node].lazy;
		}
		tree[node].lazy = 0;
	}
}
ll query(vector<Tree>& tree, int node, int s, int e, int left, int right) {
	update_lazy(tree, node, s, e);
	if (e < left || right < s)
		return 0;
	if (left <= s && e <= right)
		return tree[node].val;
	int mid = (s + e) / 2;
	return query(tree, node * 2, s, mid, left, right)
		+ query(tree, node * 2 + 1, mid + 1, e, left, right);
}
void update(vector<Tree>& tree, int node, int s, int e, 
	int left, int right, ll diff) {
	update_lazy(tree, node, s, e);
	if (e < left || right < s)
		return;
	if (left <= s && e <= right) {
		tree[node].lazy += diff;
		update_lazy(tree, node, s, e);
		return;
	}
	int mid = (s + e) / 2;
	update(tree, node * 2, s, mid, left, right, diff);
	update(tree, node * 2 + 1, mid + 1, e, left, right, diff);
	tree[node].val = tree[node * 2].val + tree[node * 2 + 1].val;
}
ll solve(int s, int e) {
	while (s + 3 <= e) {
		int p = (s * 2 + e) / 3;
		int q = (s + e * 2) / 3;
		ll A1 = query(tree, 1, 1, MAX, p, p);
		ll A2 = query(tree, 1, 1, MAX, q, q);
		ll B1 = query(cons, 1, 1, MAX, p, p);
		ll B2 = query(cons, 1, 1, MAX, q, q);
		if (A1 * p + B1 > 
			A2 * q + B2)
			s = p;
		else
			e = q;
	}
	ll ret = LNF;
	for (int i = s; i <= e; i++) {
		ll A = query(tree, 1, 1, MAX, i, i);
		ll B = query(cons, 1, 1, MAX, i, i);
		ret = min(ret, A * i + B);
	}
	return ret;
}
int main() {
	cin.tie(0)->sync_with_stdio(0);

	cin >> N >> K;
	tree.resize(1 << 18);
	cons.resize(1 << 18);
	for (int i = 1; i <= N; i++) {
		cin >> a[i];
	}
	for (int i = 1; i <= N; i++) {
		cin >> asc[i];
	}
	for (int i = 1; i <= N; i++) {
		cin >> desc[i];
	}
	for (int i = 1; i <= N; i++) {
		update(tree, 1, 1, MAX, 1, a[i]-1, -desc[i]);
		update(cons, 1, 1, MAX, 1, a[i]-1, desc[i] * a[i]);
		update(tree, 1, 1, MAX, a[i], MAX, asc[i]);
		update(cons, 1, 1, MAX, a[i], MAX, -asc[i] * a[i]);
		if (i >= K) {
			ans = min(ans, solve(1, 100000));
			update(tree, 1, 1, MAX, 1, a[i-K+1] - 1, desc[i-K+1]);
			update(cons, 1, 1, MAX, 1, a[i-K+1] - 1, -desc[i-K+1] * a[i - K+1]);
			update(tree, 1, 1, MAX, a[i-K+1], MAX, -asc[i-K+1]);
			update(cons, 1, 1, MAX, a[i-K+1], MAX, asc[i-K+1] * a[i-K+1]);
		}
	}
	cout << ans << "\n";
}

상당히 실행속도가 느리다.

삼분탐색이 아닌 이분탐색, 그리고 펜윅트리를 이용한다면 실행속도가 훨씬 빨라질 것이다.

Fenwick Tree가 재귀 segtree보다 확실히 구현량이 적어 생산성이 좋아보인다.

 

다만, 일반 최소/최대 RMQ, 구간합 segtree의 기본 펜윅과 재귀 세그먼트트리를 비교해봤을 때엔 실행속도가 거의 비슷해 당장 펜윅을 배우러갈 것 같진 않고 좀 더 필요성을 느끼면 익혀볼 듯하다.