[BOJ 2243 | C++] 사탕상자 :: 세그먼트 트리, 이분탐색

• 문제 링크 : 2243번: 사탕상자
• 문제 태그 : 자료 구조, 세그먼트 트리, 이분 탐색

 

풀이

세그먼트 트리를 사용해야 함을 알고 접근하였습니다. (사실 세그먼트 트리 문제는 모르고 접근해야 제일 좋다고 하지만..)

 

이 문제에서 가장 어려웠던 부분은 K번째로 맛있는 사탕을 찾는 일이었습니다. 

가장 먼저 리프노드에 어떤 배열이 필요한 지 생각해보았습니다. 

 

• 각 사탕마다 순위를 매긴다.

이 경우 장점은 K번째로 맛있는 사탕을 쉽게 구할 수 있습니다. 그러나 문제의 제한 조건에서 사탕의 총 개수는 2*10^9개이므로 사탕의 순위를 업데이트하는데 O(n)이 걸리는 순간 바로 시간초과입니다. 또한 배열의 크기도 10^9개 이상이라 불가능합니다. 

 

•  사탕의 맛에 따라 순위를 매긴다.

이 경우 사탕의 맛의 범위가 1부터 10^6이기 때문에 배열의 크기로 선언하기에도 적절합니다. 

 

또한 사탕의 순위는 맛이 오름차순으로 정렬(1 : 가장 맛있는 사탕, 숫자가 커질수록 맛 없음)된 상태에서 매겨지므로 순위 결정은 앞에 얼마나 더 맛있는 사탕이 있냐, 즉 사탕의 개수에 따라서 결정됩니다. 예를 들어 앞에 5개의 더 맛있는 사탕이 있다면 다음으로 맛있는 사탕은 6번째로 맛있는 사탕입니다. 

이를 통해서 prefix sum이 필요하다는 것을 알 수 있습니다. 따라서 구간합 세그먼트 트리를 구성합니다. (코드의 update 함수)

 

그리고 k번째로 맛있는 사탕을 찾아야하는데... 이때 이분탐색이 이용됩니다. 

예를 들어 4번째로 맛있는 사탕을 찾아야한다는 의미는 맛이 1인 사탕부터 순서대로 나열했을 때 4번째에 위치한 사탕을 찾아야함과 같습니다. 따라서 그림의 예시에서는 값 4는 prefix_sum의 배열에서 인덱스 3과 4 사이에 존재하며 결론적으로 4번째로 맛있는 사탕은 맛이 4인 사탕을 의미합니다.  

 

그리고 이렇게 인덱스를 찾아낼 때 시간 복잡도를 줄이기 위해 이분탐색을 이용합니다.(코드의 kth 함수) 

#include <iostream>
#include <vector>

using namespace std;
const int MAX = 1000000;

int main() {
	ios::sync_with_stdio(false);
	cin.tie(NULL);
	cout.tie(NULL);

	int n;
	cin >> n;

	int sz = 1;
	while (sz < MAX + 1) sz *= 2;
	vector tree(sz * 2, 0);

	// prefix sum
	auto update = [&](int i , int x) {
		i += sz;
		tree[i] += x;
		while (i > 1) {
			i /= 2;
			tree[i] = tree[i * 2] + tree[i * 2 + 1];
		}
	};

	auto kth = [&](int k) {
		int i = 1;
		while (i < sz) {
			if (tree[i * 2] >= k) i *= 2;
			else k -= tree[i * 2], i = 2 * i + 1;
		}
		return i - sz;
	};

	for (int i = 0;i < n;i++) {
		int a, b, c;
		cin >> a;

		if (a == 1) {
			cin >> b;
			int ans = kth(b);
			cout << ans << '\n';
			update(ans, -1);
		}
		else {
			cin >> b >> c;
			update(b, c);
		}
	}
}

 

kth 함수는 세그먼트 트리 위에서 이분탐색을 진행합니다.

인덱스는 루트 노드에서부터 시작하며, 점차 아래로 내려갑니다. 

---

ICPC Sinchon 25W 에서 중급 강사님(jinhan814)님의 세그먼트 강의를 듣고 질문을 통해 풀었습니다.🍊