[C++][백준] 사탕상자(222243) - 세그먼트트리

문제

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

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

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

 

입력

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

 

출력

A가 1인 모든 입력에 대해서, 꺼낼 사탕의 맛의 번호를 출력한다. 물론, A=2 이면서 C<0 일 때는 출력하지 않는다.

 

풀이

세그먼트트리를 활용하여 문제를 풀이하였습니다.

 

세그먼트트리의 리프노드의 경우

idx번 맛 사탕 갯수 정보를 담고있습니다.

 

리프노드가 아닌 경우 

l~r번맛 사이에 있는 사탕 갯수 정보를 담고있습니다. 

 

예를 들어 아래와 같은 상황을 가정해보겠습니다.

1. 사탕 맛 범위가 1~10입니다.

2. 테스트 케이스와 동일하게 1번맛 사탕 2개, 3번맛 사탕 3개를 넣습니다.

 

다음과 같이 세그먼트트리가 업데이트됩니다.

 

이때 소유한 사탕 중 가장 맛있는 n번째 사탕을 꺼내야합니다.

가장 맛있는 사탕은 가장 왼쪽에 있기 때문에 왼쪽 노드부터 살펴봅니다.

 

테스트 케이스와 동일하게 가장 맛있는 두번째 사탕을 꺼낸다고 가정해봅시다.

왼쪽 노드의 값이 n(2)보다 같거나 작으면 왼쪽 노드를 따라갑니다.가장 맛있는 n번째 사탕을 찾아가는 과정입니다.리프 노드에 도달시 해당 노드가 가장 맛있는 n번째 사탕이기 때문에 값을 1빼줍니다.

 

이 경우에는 두번째로 맛있는 사탕은 1번 사탕이었기 때문에 1을 출력합니다.

1

 

 

 

세그먼트 트리는 아래와 같이 업데이트 됩니다.

 

 

마찬가지로 테스트 케이스와 동일하게 가장 맛있는 두번째 사탕을 꺼내봅니다.

4번 노드를 살펴보면 왼쪽 노드의 값이 n(2)보다 같거나 작지 않기 때문에 왼쪽 노드를 따라갈 수 없습니다.

너무나 당연하게도 가장 맛있는 n번째 사탕을 찾아가야하는데

1~2번 맛 사탕의 갯수는 모두 합쳐도 n보다 작기 때문에 해당 맛 범위에는 우리가 찾으려는 가장 맛있는 n번째 사탕이 없습니다.

 

따라서 오른쪽 자식노드의 n - (왼쪽 자식 노드 value)번째 맛있는 사탕을 찾아갑니다.

 

 

 

리프 노드에 도달시 해당 노드가 가장 맛있는 n번째 사탕이기 때문에 값을 1빼줍니다.

 

이 경우에는 두번째로 맛있는 사탕은 3번 사탕이었기 때문에 3을 출력합니다.

3

 

 

 

정리해보자면

왼쪽 자식 노드가 n보다 같거나 크다면 왼쪽 노드를 따라가고

그렇지 않으면 오른쪽 노드를 따라갑니다.

리프 노드에 도달시 리프 노드를 출력합니다.

 

 

 

전체 코드는 아래와 같습니다.

#include <iostream>

using namespace std;
const int MAX = 1000010;
long long tree[MAX * 4];
int N;

// index(맛)를 value만큼 더한다.
long long update(int index, int value, int node, int l, int r) {
	if (index < l || index > r) {
		return tree[node];
	}
	else if (l == r) {
		return tree[node] += value;
	}
	int mid = (l + r) / 2;
	return tree[node] = update(index, value, node * 2, l, mid)
		+ update(index, value, node * 2 + 1, mid + 1, r);
}

// k번째 사탕의 맛을 출력한다.
long long query(int k, int node, int l, int r) {
	int mid = (l + r) / 2;
	if (l == r) {
		return l;
	}
	else if (tree[node * 2] >= k) {
		return query(k, node * 2, l, mid);
	}
	else {
		return query(k - tree[node * 2], node * 2 + 1, mid + 1, r);
	}
}

int main() {
	scanf("%d", &N);

	while (N--) {
		int instruction;
		scanf("%d", &instruction);

		if (instruction == 1) { // k번째 사탕을 구한다, k번째 사탕을 뺀다.
			int k;
			scanf("%d", &k);
			int kth_taste = query(k, 1, 1, MAX);
			cout << kth_taste << "\n"; // k번째 사탕을 구한다.
			update(kth_taste, -1, 1, 1, MAX);
		}
		else if (instruction == 2) { // 사탕을 넣거나 뺸다.
			int taste, count;
			scanf("%d %d", &taste, &count); // 맛, 갯수
			update(taste, count, 1, 1, MAX);
		}
	}
	

	return 0;
}