[BOJ 12899] 데이터 구조(c++)

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

---

문제 개요

자연수를 저장하는 데이터베이스 S에 대해 다음의 쿼리를 처리합니다.

• 유형 1 : S에 자연수 X를 추가한다
• 유형 2 : S에 포함된 숫자 중 X번째로 작은 수를 응답하고 그 수를 삭제한다.(S에 포함된 숫자는 X개 이상을 보장)

(1 ≤ N, X ≤ 2,000,000)

---

문제 풀이

이 문제는 원래 개수를 저장하는 세그먼트 트리와 이분 탐색을 통해 효율적으로 풀이 가능하지만, 바이너리 트라이로 해결하였습니다.

class node {
public:
	int cnt;
	node *child[2];
};

node pool[2000000 * 4];
node* root;
int p_cnt;

node* new_node() {
	pool[p_cnt].child[0] = nullptr;
	pool[p_cnt].child[1] = nullptr;
	return &pool[p_cnt++];
}

메모리 제한이 512MB로 매우 널널하니 메모리 풀로 최적화 하였습니다.

 

void insert(int x) {
	node* cur = root;
	for (int i = 21; i >= 0; i--) {
		cur->cnt++;
		int b;
		if (x & (1 << i)) {
			b = 1;
		}
		else {
			b = 0;
		}
		if (cur->child[b] == nullptr) {
			cur->child[b] = new_node();
		}	
		cur = cur->child[b];
	}
	cur->cnt++;
}

트라이에 값을 저장하는 함수입니다. x의 최대값이 200만이라 i를 21부터 탐색하였습니다. 몇 번째 수인지를 확인하기 위해 cnt를 넣었습니다.

int remove(int x) {
	node* cur = root;
	int ret = 0;
	for (int i = 21; i >= 0; i--) {
		cur->cnt--;
		int b;
		if (cur->child[0] && cur->child[1]) {
			if (cur->child[0]->cnt >= x) {
				b = 0;
			}
			else {
				x -= cur->child[0]->cnt;
				b = 1;
			}
		}
		else if (cur->child[0]) {
			b = 0;
		}
		else {
			b = 1;
		}
		ret |= (b << i);
		cur = cur->child[b];
	}
	cur->cnt--;
	return ret;
}

x번째 수를 찾아 제거하며 리턴하는 함수입니다.

---

코드

#include <iostream>
using namespace std;

class node {
public:
	int cnt;
	node *child[2];
};

node pool[2000000 * 4];
node* root;
int p_cnt;

node* new_node() {
	pool[p_cnt].child[0] = nullptr;
	pool[p_cnt].child[1] = nullptr;
	return &pool[p_cnt++];
}

void insert(int x) {
	node* cur = root;
	for (int i = 21; i >= 0; i--) {
		cur->cnt++;
		int b;
		if (x & (1 << i)) {
			b = 1;
		}
		else {
			b = 0;
		}
		if (cur->child[b] == nullptr) {
			cur->child[b] = new_node();
		}	
		cur = cur->child[b];
	}
	cur->cnt++;
}

int remove(int x) {
	node* cur = root;
	int ret = 0;
	for (int i = 21; i >= 0; i--) {
		cur->cnt--;
		int b;
		if (cur->child[0] && cur->child[1]) {
			if (cur->child[0]->cnt >= x) {
				b = 0;
			}
			else {
				x -= cur->child[0]->cnt;
				b = 1;
			}
		}
		else if (cur->child[0]) {
			b = 0;
		}
		else {
			b = 1;
		}
		ret |= (b << i);
		cur = cur->child[b];
	}
	cur->cnt--;
	return ret;
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);
	cout.tie(0);
	int n;
	cin >> n;
	root = new_node();
	for (int i = 0; i < n; i++) {
		int cmd, x;
		cin >> cmd >> x;
		if (cmd == 1) {
			insert(x);
		}
		else {
			int ret = remove(x);
			cout << ret << '\n';
		}
	}
	return 0;
}