[백준 2243] 사탕상자

1. 문제 

 

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

2243번: 사탕상자

---

2. 풀이

 

세그먼트 트리를 이용한 문제이다. 물론 나는 인덱스 트리로 풀었다. 인덱스 트리에 대해 잘 모른다면 먼저 공부하고 오자!

인덱스 트리 : https://kinngife.tistory.com/32

인덱스 트리 (Binary Indexed Tree)

 

이 문제에서 사탕의 맛은 1부터 100만까지의 정수로 구분되고, 입력에 따라 2가지 행동이 존재한다. 

- 1 B : 상자에서 B번째 사탕을 꺼내고 이 사탕의 맛을 출력한다. 

- 2 B C : 맛이 B인 사탕 C개를 상자에 넣는다. (C가 음수일 수도 있음)

 

B번째 사탕을 꺼낸다는 것은 사탕이 항상 정렬되어 있어야 한다는 것을 말하고, 정렬이 되어 있는 상태에서 B번째 사탕을 찾는 것은 $O(N)$만큼의 시간이 필요할 것이다. 문제는 이러한 행동을 최대 10만 번을 해야 하는데 당연히 시간 초과가 날 것이다. 결국 B번째 사탕을 찾을 때 시간을 $O(logN)$으로 줄여야 하는데, 우리는 구간 합을 구할 때 이용하는 인덱스 트리를 생각할 수 있다. 

 

사탕의 맛이 1부터 100만까지의 정수로 구분되기 때문에 인덱스 트리의 리프 노드에 사탕의 개수를 저장하면 될 것이다. 다음과 같이 초기화를 하자. 사탕이 1개도 없는 빈 상자로 시작한다고 했으므로, 모든 노드가 0인 상태로 놔두자. 또한 사탕의 총개수가 20억을 넘지 않으므로 int형 배열만으로도 충분하다. 

 

int OFFSET = 1048576;
vector<int>tree(OFFSET * 2);

 

 

먼저 사탕을 넣는 것을 생각해 보자. 맛이 B인 사탕 C개를 넣는 것은 아래와 같이 기존 update 함수를 조금만 바꾸면 된다. 

 

void update(int idx, int num) {
	idx += OFFSET;
	tree[idx] += num;
	while (idx > 1) {
		idx /= 2;
		tree[idx] = tree[idx * 2] + tree[idx * 2 + 1];
	}
}

 

기존 update 함수에서 3번째 line만 바뀐 것을 볼 수 있다. 기존 사탕에다가 사탕을 추가로 넣는 것이므로 idx번째 숫자를 num으로 update 하는 것이 아니라 idx번째 숫자에 num을 더해야 한다. 

 

 

이제 사탕을 빼는 것을 생각해 보자. 기존 구간 합을 구하는 함수는 리프 노드부터 루트 노드까지 올라가면서 값을 구하는 bottom-up 방식이었다. 하지만 이번 문제에서는 루트 노드부터 리프 노트까지 값을 구하는 top-down 방식을 이용할 것이다. 

 

인덱스 트리의 모든 노드에는 구간 합 즉, 특정 맛의 구간에 대한 사탕의 총개수가 저장되어 있을 것이다. 만약 4번째 사탕을 꺼내는 경우에 현재 노드에 5가 저장되어 있다고 해보자. 그러면 현재 노드에 저장되어 있는 5라는 값은 특정 맛의 구간에 사탕이 총 5개 있다는 뜻인데, 내가 구하려는 4번째 사탕이 포함되어 있다고 생각할 수 있으므로 자식 노드를 탐색하면 될 것이다.

반대로 내가 6번째 사탕을 꺼내는 경우에 현재 노드에 5가 저장되어 있다고 하면, 현재 노드가 저장하고 있는 특정 맛의 구간에 존재하는 사탕들에는 내가 구하려는 사탕이 없다는 뜻이므로 다른 노드를 탐색해야 한다. 그러면 어떤 노드를 탐색해야 할까? 질문을 바꿔서 6번째 사탕은 어떤 노드에 포함되어 있을까? 바로 현재 노드의 형제(sibling) 노드이다. 이유는 아래에서 좀 더 생각해 보자. 이것보다 중요한 것은 형제 노드에서 이전과 같이 6번째 사탕을 생각하는 것이 아니라 1번째 사탕을 생각해야 한다는 것이다. 왜냐하면 처음 비교한 값인 현재 노드에 5가 저장되어 있었고 만약에 현재 노드가 k번째 구간까지의 사탕을 포함한다고 했을 때 k번째 구간까지의 사탕이 5개밖에 존재하지 않는다는 뜻이고, 내가 구하려는 6번째 사탕은 (k+1) 번째 구간부터의 사탕 중 1번째 사탕이기 때문이다. 

 

정리하면 다음과 같다. 

$K$번째 사탕을 꺼낸다고 할 때 현재 노드에 저장되어 있는 숫자를 $M$이라고 하자. 

1. $K <=M$인 경우 :  현재 노드의 자식 노드 중 $K$번째 사탕을 탐색

2. $K> M$인 경우 : 현재 노드의 형제 노드 중 $K-M$번째 사탕을 탐색

 

구현을 해보자. 

루트 노드부터 리프 노드까지 top-down 방식으로 탐색을 한다고 했으므로 탐색을 할 인덱스를 1로 놓고 시작한다. 리프 노드에 도착했다는 것은 인덱스가 $OFFSET$보다 크거나 같을 때인 것으로 판단할 수 있다. 위에서 1번의 경우처럼 자식 노드를 탐색하려면 현재 인덱스에 2를 곱하여 자식 노드로 내려가면 된다. 2번의 경우처럼 형제 노드를 탐색하려면 현재 인덱스에 1을 더하면 된다. 코드는 다음과 같다. 

 

int find(int num) {
	int idx = 1;
	while (1) {
		if (idx >= OFFSET) break;

		if (tree[idx] >= num) idx *= 2;
		else {
			num -= tree[idx];
			idx++;
		}
	}

	if (tree[idx] < num) idx++;

	return idx - OFFSET + 1;
}

 

13번째 line은 리프 노드에 도착했을 때 마지막으로 값을 한 번 더 비교하여 형제 노드로 갈지 말지 판단해 주는 것이다. 탈출 조건을 $OFFSET$보다 크거나 같을 때라고 설정하였으므로 실질적으로 마지막 리프 노드에서는 값 비교를 하지 않기 때문이다. 내가 구하려는 사탕의 맛을 구하는 것이므로 반환값은 15번째 line처럼 구현했다. 

---

3. 전체 코드

 

#include <iostream>
#include <vector>
#include <algorithm>
#include <stack>
#include <queue>
#include <deque>
#include <map>
#include <set>
#include <string>
#include <climits>
#include <cmath>
#define INF INT_MAX
using namespace std;
using ll = long long;
using ii = pair<int, int>;

int OFFSET = 1048576;
vector<ll>tree(OFFSET * 2);
void update(int idx, int num) {
	idx += OFFSET;
	tree[idx] += num;
	while (idx > 1) {
		idx /= 2;
		tree[idx] = tree[idx * 2] + tree[idx * 2 + 1];
	}
}
int find(int num) {
	int idx = 1;
	while (1) {
		if (idx >= OFFSET) break;

		if (tree[idx] >= num) idx *= 2;
		else {
			num -= tree[idx];
			idx++;
		}
	}

	if (tree[idx] < num) idx++;

	return idx - OFFSET + 1;
}

int main()
{
	ios_base::sync_with_stdio(false);
	cin.tie(nullptr), cout.tie(nullptr);
	int n; cin >> n;
	while (n--) {
		int a; cin >> a;
		if (a == 1) {
			int b; cin >> b;
			cout << find(b) << "\n";
			update(find(b) - 1, -1);
		}
		else {
			int b, c; cin >> b >> c;
			update(b - 1, c);
		}
	}

	return 0;
}