[APIO 2020-A] 벽 칠하기

TMI

대회에서는 맞왜틀을 했던 문제! 구현을 좀 지저분하게 해서 맞기까지 고생을 좀 했지만(이 글을 쓰다가 깨달은 사실인데 과도한 최적화를 했기 때문에 더러워진 것이었다.) 최적화 과정에서 독특한(아님 말고) 아이디어가 사용되어 재밌었다. 이 아이디어는 다른 문제에서도 사용할 수 있을 것 같다.

---

문제 링크

boj.kr/19618

oj.uz/problem/view/APIO20_paint

---

문제 요약

N개의 구간으로 된 벽을 칠하는 데 필요한 지령의 최소 횟수를 구하는 문제.

지령(x, y): x번 일꾼이 y번째 구간을 색 C[y]로 칠하고, (x+1)%M번 일꾼이 y+1구간을 색 C[y+1]로 칠하고....이런 방식으로 M명의 일꾼이 한 구간씩 칠해, 총 M개의 구간을 칠하도록 한다. 이때 모든 일꾼에 대해, 그 일꾼이 맡은 구간의 색은 그 일꾼이 좋아하는 색 중 하나여야 된다.

---

힌트

1. DP
2. $0 \le k \le K$에 대해 $f(k) < 640$

↓2번 힌트에 대한 부가 설명

문제의 조건에서 $\sum f(k)^2 \le 400000$이므로 최악의 경우에 f(k)가 한쪽으로 다 몰빵됐다고 생각하면 시그마를 떼고 생각할 수 있고, 양변에 제곱근을 씌워주면 대충 $f(k) < 640$임을 알 수 있다.

---

풀이

편의상 문제에서 구간이라고 칭한 것을 벽으로 바꿔 말하겠다.

 

한 번 지령이 내려질 때 M개의 인접한 벽들이 색칠된다. 이때 이 길이가 M인 구간의 가장 마지막 벽(이하 'X벽' 좋은 이름이 생각이 안 나네요..ㅈㅅ)에 집중해보자.

각 벽에 대해 그 벽이 'X벽'인지를 판단할 수 있다.

(즉, 어떤 $i(M \le i \le N)$번째 벽에 대해 $[i-M+1, i]$의 벽을 M명의 사람이 순서대로 칠할 수 있는지를 판단할 수 있다.)

이 판단을 내리는 데 시간이 얼마가 걸리든 일단 구해냈다고 가정해보자.

 

그러면 이제 각 'X벽'에서 왼쪽으로 길이가 M인 선분을 뻗어내어 이 선분들이 집합이 1부터 N까지의 구간을 모두 덮는지, 또 이때 모든 구간을 덮는게 가능하다면 최소 몇개의 선분으로 모든 구간을 덮을 수 있는지를 계산하는 문제가 된다. ('공주님의 정원'같은 분위기..?)

이런 문제는 priority_queue나 세그먼트 트리를 사용해 풀 수 있고, 포인터를 점프했다가 돌아왔다가 다시 점프하고...를 반복하는 느낌으로 그리디하게 구하면 $O(N)$으로도 계산할 수도 있다.

 

자 그럼 아까 대충 넘겼던 'X벽'을 구하는 방법을 생각해보자.

가장 간단히 생각할 수 있는 방법은 $i(M \le i \le N)$번째 벽에 대해 $i-M+1$번째 벽부터 $j(0 \le j < M)$번 일꾼이 색칠을 할 수 있는지 확인하고, 그 다음 벽을 다음 일꾼이 칠할 수 있는지 확인하고....를 M번 반복하는 방식이다.

하지만 시간복잡도가 i를 정하는 데 $O(N)$, j를 정하는 데 $O(M)$, 검사가 모든 (i, j)에 대해 $O(M)$이 걸리기 때문에 $O(NM^2)$이 돼서 터진다.

시간을 어떻게 줄일 수 있을까?

 

비록 여전히 시간이 터지긴 하지만 M하나를 뗄 수 있는 방법이 있다. 아마 이렇게 하면 섭태 2~4가 해결된다. 그리고 이 접근이 100점 풀이와 연결되니 주의깊게 봐보자.

N*M 사이즈의 DP 테이블을 잡아주자. 이때 dp[i][j] = i번째 벽을 j번째 일꾼이 칠할 때, 최대 몇명 연속으로 색칠을 하고 있는가로 정의하자.

그러면 dp[i][j]의 값은 $O(1)$에 다음과 같이 계산할 수 있다.

• i번째 벽을 j번째 일꾼이 칠할 수 있을 때(=C[i]가 j번째 일꾼이 좋아하는 색에 포함될 때): dp[i][j] = dp[i-1][(j-1+M)%M]+1
• i번째 벽을 j번째 일꾼이 칠할 수 없을 때: dp[i][j] = 0

이때, 어떤 $j(0 \le j < M)$에 대해 $d[i][j] \ge M$이면 i번째 벽은 'X벽'이다. ('X벽'과 dp[i][j]의 정의를 생각해보면 자명하다.)

 

참, i번째 벽을 j번째 일꾼이 칠할 수 있는지를 어떻게 빠르게 알 수 있냐면 입력으로 주어지는 B배열을 전처리 해두면 된다.

인덱스가 값이 되고, 값이 인덱스가 되는 느낌으로 '색깔 B[i][j]를 i가 좋아한다'라고 저장하면 각 색에 대해 그 색으로 칠할 수 있는 일꾼의 번호가 오름차순으로 정렬된다. 어차피 각 벽에 대해 dp값을 계산할 때 일꾼의 번호가 증가하는 순서대로 값을 채우기 때문에 투포인터를 이용해 검사해주면 빠르게 확인이 가능하다.

 

이렇게 DP를 이용해 'X벽'을 조금 빠르게 구할 수 있었다. 여기서 시간을 더 줄이려면 두번째 힌트인 $0 \le k \le K$에 대해 $f(k) < 640$임을 이용해 DP를 최적화 해주면 된다.

$f(k) < 640$라는 말은 각 벽에 대해 색칠을 할 수 있는 일꾼의 수가 640명보다 작다는 것이고, 이는 각 벽에 대해 최대 640개의 dp값이 쓸모있는 정보가 된다는 것이다. (나머지는 전부 0이라고 생각하면 되니까)

이렇게 각 벽에 대해 최대 640개씩만 계산하게 되면 시간복잡도가 O(640N)으로 시간제한 안으로 들어오는 것을 볼 수 있다.

 

하지만 여기까지만 하면 메모리가 터져서 DP 테이블도 최적화 해줬다.

DP식을 보면 dp[i] 줄을 계산할 때 dp[i-1] 줄의 정보만 이용하기 때문에 i-2이하의 줄은 굳이 보관할 필요가 없어진다.

즉, $O(640*2)$ 정도의 공간만으로도 dp를 계산할 수 있다. 하지만 이렇게 구현하려면 조금 지저분해지니 그냥 $O(2M)$을 잡고 빈칸을 많이 남긴채로 써도 될 것 같다.

이런 방법의 공간 최적화는 '헤르메스' 맞나? 등의 문제에서 자주 볼 수 있다.

 

 

앞에서도 말했지만 구현할 때 공간복잡도에 대해 깊게 고민하지 않고 막 짜서 많이 지저분하고, 심지어 나중에 최적화를 덕지덕지 붙이면서 과도한 최적화가 되어버렸다. 아무튼 더럽게 짰다;

#include "paint.h"
#include <vector>
#include <iostream>
#define N 100005
using namespace std;

int ba, sg[N * 4];
bool v[N];
vector<int> di, dc, ti, tc, ve[N];

void up(int p, int q) {
	p += ba;
	sg[p] = q;
	for (p /= 2; p > 0; p /= 2) {
		sg[p] = min(sg[p * 2], sg[p * 2 + 1]);
	}
}

int get(int p, int q) {
	int re = 1e9;
	p += ba; q += ba;
	while (p < q) {
		if (p % 2 == 1) re = min(re, sg[p]), p++;
		if (q % 2 == 0) re = min(re, sg[q]), q--;
		p /= 2; q /= 2;
	}
	if (p == q) re = min(re, sg[p]);
	return re;
}

int minimumInstructions(int n, int m, int K, vector<int> c, vector<int> a, vector<vector<int> > b) {
	int i, j, k;

	c.push_back(0);
	for (i = c.size() - 1; i > 0; i--) {
		c[i] = c[i - 1];
	}

	for (i = 0; i < m; i++) {
		for (j = 0; j < a[i]; j++) {
			ve[b[i][j]].push_back(i);
		}
	}
	di.push_back(1e9);
	dc.push_back(-1e9);
	for (i = 1; i <= n; i++) {
		ti.clear(); tc.clear();
		for (j = 0; j < ve[c[i]].size(); j++) {
			ti.push_back(ve[c[i]][j]);
			tc.push_back(0);
		}
		ti.push_back(1e9);
		tc.push_back(-1e9);
		for (j = k = 0; j < ti.size() - 1; j++) {
			if (ti[j] == 0) {
				if (di.size() > 1 && di[di.size() - 2] == m - 1) {
					tc[j] = dc[di.size() - 2] + 1;
				} else {
					tc[j] = 1;
				}
			} else {
				for (; di[k] < ti[j] - 1; k++);
				if (di[k] == ti[j] - 1) {
					tc[j] = dc[k] + 1;
				} else {
					tc[j] = 1;
				}
			}
		}
		di = ti;
		dc = tc;
		for (j = 0; j < dc.size() - 1; j++) {
			if (dc[j] >= m) {v[i] = 1; break;}
		}
	}
	for (ba = 1; ba < n + 1; ba *= 2);
	for (i = 0; i < ba * 2; i++) sg[i] = 1e9;
	up(0, 0);
	for (i = m; i <= n; i++) {
		if (v[i]) {
			int t = get(i - m, i - 1) + 1;
			if (t < sg[ba + i]) up(i, t);
		}
	}
	if (get(n, n) < 1e9) return get(n, n);
	else return -1;
}