[백준] 2243번: 사탕상자 (c++)

 

문제

---

수정이는 어린 동생을 달래기 위해서 사탕을 사용한다. 수정이는 평소에 여러 개의 사탕을 사서 사탕상자에 넣어두고, 동생이 말을 잘 들을 때면 그 안에서 사탕을 꺼내서 주곤 한다.

각각의 사탕은 그 맛의 좋고 나쁨이 1부터 1,000,000까지의 정수로 구분된다. 1이 가장 맛있는 사탕을 의미하며, 1,000,000은 가장 맛없는 사탕을 의미한다. 수정이는 동생이 말을 잘 들은 정도에 따라서, 사탕상자 안에 있는 사탕들 중 몇 번째로 맛있는 사탕을 꺼내주곤 한다. 예를 들어 말을 매우 잘 들었을 때에는 사탕상자에서 가장 맛있는 사탕을 꺼내주고, 말을 조금 잘 들었을 때에는 사탕상자에서 여섯 번째로 맛있는 사탕을 꺼내주는 식이다.

수정이가 보관하고 있는 사탕은 매우 많기 때문에 매번 사탕상자를 뒤져서 꺼낼 사탕을 골라내는 일은 매우 어렵다. 수정이를 도와주는 프로그램을 작성하시오.

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

2243번: 사탕상자

 

문제조건

---

첫째 줄에 수정이가 사탕상자에 손을 댄 횟수 n(1 ≤ n ≤ 100,000)이 주어진다. 다음 n개의 줄에는 두 정수 A, B, 혹은 세 정수 A, B, C가 주어진다. A가 1인 경우는 사탕상자에서 사탕을 꺼내는 경우이다. 이때에는 한 정수만 주어지며, B는 꺼낼 사탕의 순위를 의미한다. 이 경우 사탕상자에서 한 개의 사탕이 꺼내지게 된다. 또, A가 2인 경우는 사탕을 넣는 경우이다. 이때에는 두 정수가 주어지는데, B는 넣을 사탕의 맛을 나타내는 정수이고 C는 그러한 사탕의 개수이다. C가 양수일 경우에는 사탕을 넣는 경우이고, 음수일 경우에는 빼는 경우이다. 맨 처음에는 빈 사탕상자에서 시작한다고 가정하며, 사탕의 총 개수는 2,000,000,000을 넘지 않는다. 또한 없는 사탕을 꺼내는 경우와 같은 잘못된 입력은 주어지지 않는다.

 

풀이

---

인덱스 트리를 사용하여 풀어주었다. 인덱스 트리의 노드는 각 가중치를 가진 사탕의 개수를 의미한다. 예를 들어 4종류의 사탕이 있고, 각 개수가 2개, 3개, 1개, 1개가 있다고 한다면 인덱스 트리는 아래 그림과 같게 된다.

 

여기서 4순위의 사탕을 찾으려고 한다면 아래 그림과 같은 로직으로 찾을 수 있다.

1️⃣ 노드의 값이 찾으려는 순위보다 크다면 왼쪽 트리를 본다
2️⃣ 노드의 값이 찾으려는 순위보다 작다면 오른쪽 트리를 본다. 단, 왼쪽 트리의 순위를 지나친 것임으로 구하려는 순위를 업데이트해준다.
3️⃣ 리프노드에 도착했다면 답

 

코드

---

#include <iostream>
#define MAX 1000000
 
using namespace std;
int tree[4 * MAX + 2];
int temp = 1;
 
void update(int idx, int c) {
	while (idx > 0) {
		tree[idx] += c;
		idx /= 2;
	}
}
 
int getCandy(int b) {
	int idx = 1;
	while (1) {
		if (tree[idx] >= b) {
			// leaf라면
			if (idx >= temp) {
				update(idx, -1);
				return idx - temp + 1;
			}
			idx *= 2;
			continue;
		}
		else {
			// 오른쪽 트리
			b = b - tree[idx];
			idx++;
			continue;
		}
	}
}
 
int main() {
	ios::sync_with_stdio;
	cin.tie(0);
 
	int n;
	cin >> n;
 
	while (temp < MAX) {
		temp *= 2;
	}
 
	for (int i = 0; i < n; i++) {
		int a, b, c;
		cin >> a >> b;
		if (a == 2) {
			cin >> c;
			update(b+temp-1, c);
		}
		else {
			cout << getCandy(b) << "\n";
		}
	}
}