[BAEKJOON] 2243 사탕상자 (C++)

문제링크

2243번: 사탕상자 (acmicpc.net)

2243번: 사탕상자

 

문제

요약하자면 번호가 매겨진 사탕을 상자에 일정 수만큼 넣거나 뺄 수 있다.

상자 안에서 k번째로 큰 사탕의 번호가 무엇인지 출력해야한다. 

위의 두 명령이 총 100000번까지 실행될 수 있다.

 

풀이

사탕을 넣고 뺄 때마다 정렬하고 k번째에 있는 수를 출력하면 편하겠지만, 시간초과가 날 것이다.

사탕을 계속해서 넣고 뺄 수 있고, 그럴 때마다 사탕상자의 상태가 바뀐다. 세그먼트 트리를 이용하면 시간초과문제를 해결할 수 있다.

세그먼트트리는 각 구간에 포함된 사탕의 수를 가지고 있다. 

초기에 1부터 N구간에서 k번째 수가 무엇인지 찾아야한다.

이때 중간지점  mid = (1+N)/2  을 기준으로 

k > mid 이라면 오른쪽 구간에서 k- mid 번째 수를 찾아 return 하면 된다.

k<=mid 라면 왼쪽 구간에서 k번째 수를 찾아 return 한다.

각 query마다  탐색구간이 절반씩 줄어들기 때문에 한 번 탐색을 log N번 안에 수행할 수 있다.

 

코드

#include <bits/stdc++.h>

#define endl "\n"

typedef long long ll;
using namespace std;

int N;
ll tree[4040404];

void update(int n, int s, int e, int val, int cnt) {

	if (s == e) {
		tree[n] += cnt;
		return;
	}
	int mid = (s + e) / 2;
	if(mid>=val)update(n * 2, s, mid, val, cnt);
	else update(n * 2 + 1, mid + 1, e, val, cnt);
	tree[n] = tree[n * 2] + tree[n * 2 + 1];
}

int query(int n, int s, int e, int val) {
	int rt = 0;
	int mid = (s + e) / 2;
	if (s == e) { tree[n]--; return s; }
	if (tree[n * 2] < val) rt = query(n * 2 + 1, mid+1, e, val - tree[n * 2]);
	else rt = query(n * 2, s, mid, val);
	tree[n] = tree[n * 2] + tree[n * 2 + 1];
	return rt;
}
void Solve() {
	cin >> N;
	int a, b, c;
	while (N-- > 0) {
		cin >> a;
		if (a == 1) {
			cin >> b;
			cout<<query(1, 1, 1000000, b)<<endl;
		}
		else {
			cin >> b >> c;
			update(1, 1, 1000000, b, c);
		}
	}
}

int main() {
	ios::sync_with_stdio(false); cin.tie(NULL);
	Solve();
	return 0;
}