백준2243 사탕상자

작성일 | In

문제 링크

  • http://icpc.me/2243

사용 알고리즘

  • Segment Tree

시간복잡도

  • O(N log 1e6)

풀이

간단히 요약하자면, 업데이트 쿼리가 있을 때 k번째 숫자를 구하는 문제입니다.

먼저 합을 저장하는 세그먼트 트리를 만들어주고, i번째 리프노드에 사탕의 맛이 i인 사탕의 개수를 저장해줍시다.
이렇게 되면 사탕을 넣거나 빼는 쿼리, 즉 업데이트는 쉽게 처리할 수 있겠죠.

이제 k번째 사탕을 구하는 함수를 구현해봅시다.
사탕의 개수는 리프노드(단말노드)에 저장했기 때문에 무조건 리프노드까지 가야합니다.
리프노드가 아닌 노드, 즉 내부노드를 보고 있다고 합시다. 리프노드까지 가기 위해서는 왼쪽 자식으로 갈지 오른쪽 자식으로 갈지 정해야 합니다.
이진 탐색 비슷한 느낌으로 k <= tree[node*2]라면 왼쪽 자식으로 이동하고, 그렇지 않다면 오른쪽 자식으로 이동하면 됩니다.
아래와 같은 느낌으로 구현하면 됩니다.

1
2
3
4
5
6
int query(int n, int s, int e, int k){
	if (s == e) return s;
	int m = s + e >> 1;
	if (k <= tree[n << 1]) return query(n << 1, s, m, k);
	else return query(n << 1 | 1, m + 1, e, k - tree[n << 1]);
}

전체 코드

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
#include <iostream>
#include <vector>
using namespace std;

int tree[4040404];
int bias = 1;

void update(int x, int v){
	x |= bias;
	x--;
	tree[x] += v;
	while (x >>= 1){
		tree[x] = tree[x << 1] + tree[x << 1 | 1];
	}
}

int query(int n, int s, int e, int k){
	if (s == e) return s;
	int m = s + e >> 1;
	if (k <= tree[n << 1]) return query(n << 1, s, m, k);
	else return query(n << 1 | 1, m + 1, e, k - tree[n << 1]);
}

int main(){
	ios_base::sync_with_stdio(0); cin.tie(0);
	int n; cin >> n;
	for (; bias < 1000000; bias <<= 1);
	while (n--){
		int a, b, c; cin >> a >> b;
		if (a == 1){
			int idx = query(1, 1, bias, b);
			cout << idx << "\n";
			update(idx, -1);
		}
		else{
			cin >> c;
			update(b, c);
		}
	}
}