백준2243(사탕상자)

jh Seo·2023년 4월 27일

백준 세그먼트 트리 한번에 푼 문제

백준

목록 보기

153/194

개요

백준2243(사탕 상자)

입력
첫째 줄에 수정이가 사탕상자에 손을 댄 횟수 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인 모든 입력에 대해서, 꺼낼 사탕의 맛의 번호를 출력한다.

접근방식

이분탐색을 통해 사탕의 순위를 탐색하는 방법

한 컨테이너에서 뺐다 넣었다 하면서 순위를 매기고, 몇개인지 확인한다는 조건과 100만의 범위를 보고 세그먼트 트리를 떠올렸다.
일단 세그먼트 트리의 리프노드의 index를 맛의 순위로 정하고,
리프노드의 value를 index 맛 사탕의 갯수로 정했다.
고민을 했던 부분은 A가 1일때는 맛의 순위가 아닌 사탕의 순위를
기준으로 사탕을 꺼내야 하는 부분이다.
따라서 사탕의 순위를 index로 하고 value에 맛의 순위를 저장할 컨테이너를 만들려고 봤더니 사탕이 최대 20억개라 포기했다.
따라서 기존 설정한 세그먼트 트리를 유지하며 사탕의 순위를 어떻게 판별할까 고민하다가, 처음 떠올린 방법은 0번째부터 100만번째 맛까지 인덱스를 1씩 증가하며 구간합이 찾으려는 사탕순위보다 커질때까지 반복하려 했다.
비효율적이라서 이분탐색을 이용하였다.

//targetAmount 원하는 사탕 순위보다 구간합이 처음으로 커지는 index를 찾아서 반환
int binarySearch(int targetAmount) {
//low 랑 high랑 1차이 나게 하려고 하다보니 index 0이 답일때 0을 반환못한다.

	int low=0,high=1'000'000,mid = 0;
	if (segTree[firstLeafNodeIdx] >= targetAmount) return 0;
	while (low + 1 < high) {
		mid = (low + high) / 2;
		if (HowManyCandiesInTargetRange(0, mid, 1, 0, firstLeafNodeIdx - 1) < targetAmount) low = mid;
		else high = mid;
	}
	return high;
}

세그먼트 트리의 루트 노드부터 자식 비교하면서 순회하여 사탕 순위 찾는 방법

위 방식 처럼 HowManyCandiesInTargetRagne함수를 통해 구간의 합을 일일히 구하지 말고 루트노드부터 탐색하여 원하는 사탕순위 값보다 작거나 같은 구간합을 구하는 방식이다.

루트노드부터 시작해서 인자로 들어온 targetAmount보다 작거나 같은
값을 찾아 재귀한다.

//targetAmount 원하는 사탕 순위보다 구간합이 처음으로 커지는 index를 찾아서 반환
int SearchTasteIndex(int idx,int targetAmount) {
	if (idx >= firstLeafNodeIdx) return idx - firstLeafNodeIdx;
	int tmp = segTree[idx*2];
	//왼쪽자식의 노드값이 원하는 타겟보다 작거나 같으면 왼쪽자식으로 진행
	if (targetAmount <= tmp) return SearchTasteIndex(idx * 2, targetAmount);
	//오른쪽 자식 노드값이 원하는 타겟보다 크면 오른쪽 자식으로 진행
	//오른쪽으로 꺾으니 targetAmount 에서 왼쪽자식구간합인 tmp값을 빼줘야 원하는 구간합 탐색이 가능하다.
	else return SearchTasteIndex(idx * 2 + 1, targetAmount-tmp);
}

주의 할점은 오른쪽 자식으로 재귀를 진행할때다.
오른쪽 자식은 오른쪽 구간의 합만 저장되어있으므로 targetAmount값에서
왼쪽 자식구간 합을 빼줘야한다.

이분탐색으로 순위 탐색하는 코드

#include<iostream>

using namespace std;
//2^21
//index는 맛의 순위, value는 갯수
int segTree[2097152];
int N, firstLeafNodeIdx=2097152/2;

//[targetL, targetR]맛 구간의 사탕갯수의 합을 반환하는 함수
int HowManyCandiesInTargetRange(int targetL, int targetR, int nodeNum, int curL, int curR) {
	if (targetL > curR || curL > targetR) return 0;
	if (targetL <= curL && curR <= targetR) return segTree[nodeNum];
	int mid = (curL + curR) / 2;
	return HowManyCandiesInTargetRange(targetL, targetR, nodeNum * 2, curL, mid) +
		HowManyCandiesInTargetRange(targetL, targetR, nodeNum*2+1,mid+1,curR);
}

//세그트리의 idx+firstLeafNodeIdx 인덱스의 값을 변경 후 조상노드들도 변경 
void SetAncestorNode(int idx,int val) {
	int tmpIdx = idx + firstLeafNodeIdx;
	while (tmpIdx > 0) {
		segTree[tmpIdx] += val;
		tmpIdx /= 2;
	}
}

//targetAmount 원하는 사탕 순위보다 구간합이 처음으로 커지는 index를 찾아서 반환
int binarySearch(int targetAmount) {
	int low=0,high=1'000'000,mid = 0;
	//low 랑 high랑 1차이 나게 하려고 하다보니 index 0이 답일때 0을 반환못한다.
	if (segTree[firstLeafNodeIdx] >= targetAmount) return 0;
	while (low + 1 < high) {
		mid = (low + high) / 2;
		if (HowManyCandiesInTargetRange(0, mid, 1, 0, firstLeafNodeIdx - 1) < targetAmount) low = mid;
		else high = mid;
	}
	return high;
}

void Solution(int A, int B, int C=0) {

	if (A == 1) {
		int tmpIdx = binarySearch(B);
		//index는 0부터지만 사탕의 맛은 1부터
		cout << tmpIdx+1<<'\n';
		SetAncestorNode(tmpIdx,-1);
	}
	else {
		SetAncestorNode(B,C);
	}
}
void Input() {
	int A=0, B=0, C = 0;
	cin >> N;
	for (int i = 0; i < N; i++) {
		cin >> A;
		if (A == 1) {
			cin >> B;
			Solution(A, B);
		}
		else {
			cin >> B >> C;
			Solution(A, B - 1, C);
		}
	}
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	Input();
}

세그먼트 트리에서 순회하며 순위탐색하는 코드

#include<iostream>

using namespace std;
//2^21
//index는 맛의 순위, value는 갯수
int segTree[2097152];
int N, firstLeafNodeIdx=2097152/2;

//세그트리의 idx+firstLeafNodeIdx 인덱스의 값을 변경 후 조상노드들도 변경 
void SetAncestorNode(int idx,int val) {
	int tmpIdx = idx + firstLeafNodeIdx;
	while (tmpIdx > 0) {
		segTree[tmpIdx] += val;
		tmpIdx /= 2;
	}
}

//targetAmount 원하는 사탕 순위보다 구간합이 처음으로 커지는 index를 찾아서 반환
int SearchTasteIndex(int idx,int targetAmount) {
	if (idx >= firstLeafNodeIdx) return idx - firstLeafNodeIdx;
	int tmp = segTree[idx*2];
	//왼쪽자식의 노드값이 원하는 타겟보다 작거나 같으면 왼쪽자식으로 진행
	if (targetAmount <= tmp) return SearchTasteIndex(idx * 2, targetAmount);
	//오른쪽 자식 노드값이 원하는 타겟보다 크면 오른쪽 자식으로 진행
	//오른쪽으로 꺾으니 targetAmount 에서 왼쪽자식구간합인 tmp값을 빼줘야 원하는 구간합 탐색이 가능하다.
	else return SearchTasteIndex(idx * 2 + 1, targetAmount-tmp);
}

void Solution(int A, int B, int C=0) {

	if (A == 1) {
		//B개의 인덱스 찾기 
		int tmpIdx = SearchTasteIndex(1,B);
		//index는 0부터지만 사탕의 맛은 1부터
		cout << tmpIdx+1<<'\n';
		SetAncestorNode(tmpIdx, -1);
	}
	else {
		SetAncestorNode(B,C);
	}
}
void Input() {
	int A=0, B=0, C = 0;
	cin >> N;
	for (int i = 0; i < N; i++) {
		cin >> A;
		if (A == 1) {
			cin >> B;
			Solution(A, B);
		}
		else {
			cin >> B >> C;
			Solution(A, B - 1, C);
		}
	}
}

int main() {
	ios_base::sync_with_stdio(0);
	cin.tie(0);
	Input();
}

문풀후생

세그먼트트리에 해당 갯수를 이분탐색을 통해 탐색하는 부분을 구현할때 재밌었다.

세그먼트 트리에서 루트노드부터 자식 노드들로 탐색할 때,
처음에

//targetAmount 원하는 사탕 순위보다 구간합이 처음으로 커지는 index를 찾아서 반환
int SearchTasteIndex(int idx,int targetAmount) {
	if (idx >= firstLeafNodeIdx) return idx - firstLeafNodeIdx;
	int tmp = segTree[idx*2];
	//왼쪽자식의 노드값이 원하는 타겟보다 작으면 왼쪽자식으로 진행
	if (targetAmount < tmp) return SearchTasteIndex(idx * 2, targetAmount);
	//오른쪽 자식 노드값이 원하는 타겟보다 크거나 같으면 오른쪽 자식으로 진행
	//오른쪽으로 꺾으니 targetAmount 에서 왼쪽자식구간합인 tmp값을 빼줘야 원하는 구간합 탐색이 가능하다.
	else return SearchTasteIndex(idx * 2 + 1, targetAmount-tmp);
}

이런식으로 원하는 타겟값보다 크거나 같으면 오른쪽 자식으로 탐색하도록 짜고 함수호출을
동일하게 SearchTasteIndex(1,B);로 호출했더니 틀렸다.
고민해본 결과, 그 이유는 tmp값과 targetAmount값이 같다면
현재 인덱스의 왼쪽 자식의 값이 답이라는 뜻인데,
코드구현대로 진행하면 오른쪽 자식으로 재귀가 진행되어서
오른쪽자식의 제일 왼쪽 리프노드로 진행하게되고 답을 이상하게 내기때문이였다.

따라서 위의 함수대로 짜면 SearchTasteIndex(1,B-1) 이렇게 1뺀 값을 호출하던가,

//targetAmount 원하는 사탕 순위보다 구간합이 처음으로 커지는 index를 찾아서 반환
int SearchTasteIndex(int idx,int targetAmount) {
	if (idx >= firstLeafNodeIdx) return idx - firstLeafNodeIdx;
	int tmp = segTree[idx*2];
	//왼쪽자식의 노드값이 원하는 타겟보다 작으면 왼쪽자식으로 진행
	if (targetAmount <= tmp) return SearchTasteIndex(idx * 2, targetAmount);
	//오른쪽 자식 노드값이 원하는 타겟보다 크거나 같으면 오른쪽 자식으로 진행
	//오른쪽으로 꺾으니 targetAmount 에서 왼쪽자식구간합인 tmp값을 빼줘야 원하는 구간합 탐색이 가능하다.
	else return SearchTasteIndex(idx * 2 + 1, targetAmount-tmp);
}

이런식으로 tmp값과 targetAmount값이 같으면 왼쪽 자식으로 재귀가 되도록 변경해야한다.

jh Seo

코딩 창고!

이전 포스트

백준 9345(디지털 비디오 디스크)

다음 포스트