BOJ : 2243 사탕상자

---

[BOJ 2243(P5)리뷰]

세그먼트 트리를 응용하는 문제다.

사탕상자에서 사탕을 넣거나 뺄 수 있고, 동생에게 N번째로 맛있는 사탕을 줘야할 때도 있다.

사탕의 맛은 1부터 100만이다. 이걸 배열로 기록하여 값을 업데이트 한다면 시간복잡도O(N^2)으로 시간초가가 나게 된다. 따라서 업데이트를 O(logN)에 가능하게 하여 최종적으로 O(NlogN)으로 풀어야 정답처리를 받을 수 있다.

일반적으로 세그먼트 트리에는 구간의 합을 저장한다. 가장 리프노트에는 본래의 값이 들어있고, 리프노드에서 부모노드로 올라갈 수록 자식노드의 합들을 저장하며 루트노드까지 올라간다. 최종적으로 루트노드에는 전체범위의 합이 저장되어 있다.

이 문제는 이러한 구조를 갖는 세그먼트 트리를 응용하여 풀어야 하는 문제다. 사탕의 맛은 1부터 100만까지 존재하니 1~1000000 범위의 세그먼트 트리를 만들자. 리프노드의 1번노드에는 1맛 사탕의 갯수가 저장 될 것이고 리프노드의 100만번에는 100만 맛 사탕의 갯수가 저장 될 것이다.

이해를 쉽게하기 위하여 총 8가지 맛의 사탕이 있고 1번맛 2개 4번맛 3개 8번맛 1개의 사탕이 있다고 해보자.

세그먼트 트리를 만들게 되면 다음과 같은 모양을 가질 것이다.

이 때 만약에 2번째로 맛있는 사탕을 달라는 요청이 오면 어떻게 해야할까?

먼저 루트노드에 4번째로 맛있는 사탕을 달라고 요청을 한다.
루트노드의 범위에 포함되어 있으므로 자신의 자식노드에 동일하게 요청을 한다.
이 때, 왼쪽노드는 그대로 요청을 하면 되지만 오른쪽노드에는 왼쪽노드의 값을 뺀 값을 요청해야 한다.
오른쪽 노드 기준으로 왼쪽 노드는 자신보다 앞에 위치한 값들이기 때문이다.

예를들어 6번째로 맛있는 사탕을 달라고 할때, 루트노드에서 왼쪽노드에겐 6을 요청, 오른쪽노드에겐 1을 요청한다. 이미 자신보다 앞에 자신보다 맛있는 사탕이 5개가 존재하기 때문이다. (1이 가장 맛있는 사탕이기에 가능하다.)

이런식으로 사탕이 자신의 노드 범위에 속하는지 확인을 하고 쿼리를 아래로 요청하는 작업을 반복하다가 리프노드에 도착하면 해당하는 노드의 인덱스를 반환하면된다. 노드를 내려갈때는 내려갈때마다 자신이 범위에 속해있다면 1을 감소시켜야 한다.(사탕은 1개만 주므로)

말로 설명하니 복잡하지만 루트노드로 부터 시작하여 아래로 쿼리를 전달할 때 자신이 사탕을 꺼내줄 수 없는 범위에 속한 노드라면 굳이 아래로 쿼리를 보내지 않는다는 점에서 시간 절약이 가능하다. 즉, 사탕을 꺼내주기 위하여 모든 노드를 확인할 필요가 없다.

/*
	21.01.14
	BOJ : 2243 사탕상자 (https://www.acmicpc.net/problem/2243)
	자료구조/세그먼트 트리
*/
#include <bits/stdc++.h>
using namespace std;

#define ll long long int
ll* nums;
ll* tree;
int N, Q, S;

#define MAX 1000000

void init() {
	for (int i = 1; i < 2 * S; i++) {
		tree[i] = 0;
	}
}

ll query(int node, int start, int end, int candy) {	
	if (candy > tree[node] || candy < 1) return 0;

	tree[node] -= 1;
	if (start == end) {
		return start;
	}
	int mid = (start + end) / 2;
	int d = tree[node * 2];
	return query(node * 2, start, mid, candy) + query(node * 2 + 1, mid + 1, end, candy - d);
}

void update(int node, int left, int right, int index, ll diff) {
	if (left <= index && index <= right) {
		tree[node] += diff;
		if (left != right) {
			int mid = (left + right) / 2;
			update(node * 2, left, mid, index, diff);
			update(node * 2 + 1, mid + 1, right, index, diff);
		}
	}
}


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

	S = ceil(log2(MAX));
	S = (1 << S);

	tree = new ll[2 * S];

	init();

	cin >> N;
	for (int i = 0; i < N; i++) {
		int A, B, C;
		cin >> A;
		if (A == 1) {
			cin >> B;
			cout << query(1, 1, S, B) <<'\n';
		}
		else if (A == 2) {
			cin >> B >> C;
			update(1, 1, S, B, C);
		}
	}

	return 0;
}