[P5] 백준 2243번 사탕상자 C++ 세그먼트 트리, 이분 탐색

리뷰

 

세그먼트 트리의 기초가 되는 문제인듯 싶다, 구간 업데이트 쿼리가 아닌 리프 노드까지 도달해야 하는 문제

수열과 쿼리 시리즈 부터 풀다보니 이런 기초 세그를 응용하는 문제가 약해서 생각의 폭을 넓혀야 겠다 ㅠ

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

 

전역 변수

• tree : 세그먼트 트리 정보를 저장할 정수형 배열, 사탕의 맛 개수가 최대 100만이므로 400만 이상의 크기로 세팅한다.

 

함수

1. update

void update(int node, int start, int end, int idx, int val)

 

1. 매개변수 idx의 맛을 가진 사탕의 개수를 val만큼 더해준다.
2. val에는 음수가 올 수 있으나 문제에 없는 사탕을 꺼내는 경우와 같은 잘못된 입력은 주어지지 않는다. 라고 명시되어 있다.
3. 사탕을 맛있는 순서대로 누적합을 통해 세그먼트 트리를 관리해 준다.

2. query

int query(int node, int start, int end, int idx)

 

1. 루트부터 시작하여 각 왼쪽 오른쪽 노드의 사탕 개수를 세어 사탕의 순위를 찾는 함수
2. 만약 10번째 사탕의 맛을 찾고 싶은데 루트의 왼쪽이 11이고 오른쪽이 29라면 왼쪽으로 이동한다.
3. 루트의 왼쪽이 9이고 오른쪽이 11이라면 오른쪽으로 이동하고 찾고자 하는 값에서 왼쪽 9만큼을 빼준다.
4. 그럼 이제 오른쪽으로 이동한 상태에서 1번째 사탕의 맛을 찾아주면 된다.
5. 오른쪽 노드로 이동할때의 연산을 떠오르는 것이 이 문제의 킥인 것 같다.

 

문제풀이

1. 사탕이 없는 상태로 시작하므로 수열을 입력 받을 필요도 없고 세그먼트 트리를 build해줄 필요가 없다.
2. n값을 입력 받고 n개의 쿼리를 수행해 준다, order에 따라 update와 query가 나누어 진다.
3. order가 2일 경우 index, val을 추가로 입력 받고 index맛의 사탕을 val개 추가해 준 후 상위 노드를 갱신한다.
4. 1일 경우 index를 추가로 입력 받고 query함수를 사용해 index번째 맛있는 사탕의 맛을 result에 저장한다.
5. result를 출력해 준 후에 update함수를 통해 맛이 result인 사탕의 개수를 1만큼 줄여준다.

 

참고 사항

1. 세그먼트 트리의 build가 필요 없고, tree는 0인상태에서 시작한다.
2. 업데이트 시 음수가 입력될 수도 있다, 다만 보유한 사탕 개수가 음수가 되는 케이스는 없다.
3. 탐색을 할때 오른쪽으로 가는 경우 찾고자 하는 값을 왼쪽 사탕의 개수만큼을 빼주어야 한다.
4. 입력은 1-based-indexing으로 주어진다.

 

정답 코드

#include<iostream>

using namespace std;

int tree[4000004];

void update(int node, int start, int end, int idx, int val) {
	if (idx < start || end < idx) return;
	if (start == idx && end == idx) {
		tree[node] += val;
		return;
	}
	int mid = (start + end) / 2;
	update(node * 2, start, mid, idx, val);
	update(node * 2 + 1, mid + 1, end, idx, val);
	tree[node] = tree[node * 2] + tree[node * 2 + 1];
}

int query(int node, int start, int end, int idx) {
	if (start == end) return start;
	int mid = (start + end) / 2;
	if (tree[node * 2] >= idx) return query(node * 2, start, mid, idx);
	return query(node * 2 + 1, mid + 1, end, idx - tree[node * 2]);
}

int main() {
	ios::sync_with_stdio(0);
	cin.tie(0);

	int n; cin >> n;
	while (n--) {
		int order; cin >> order;
		if (order == 2) {
			int index, val; cin >> index >> val;
			update(1, 1, 1000000, index, val);
		}
		else {
			int index; cin >> index;
			int result = query(1, 1, 1000000, index);
			cout << result << "\n";
			update(1, 1, 1000000, result, -1);
		}
	}
	
}