c++/알고리즘

[c++] 백준 사탕상자 - 인덱스드 트리

hojung 2023. 2. 2.

728x90

1. 문제

2. 입 출력

예제 입 출력

4. 문제 해설

문제의 조건을 살펴보자

사탕 상자에 수정이가 손을 댄 횟수는 n(1 <= n <= 100000)이다. 또한 사탕은 1000000 맛까지 존재한다고 한다. 만약 완전 탐색으로 이 문제를 접근한다고 한다면 최악의 경우

1000000 * 100000 = 100000000000의 연산횟수를 가지게 된다. 일반적으로 c++언어에서 1초안에 결과가 나오려면 연산횟수 1억번 안에 답이 나와야 한다. 하지만 이 문제는 O(MN) 천억번의 연산 횟수를 가지므로 완전탐색으로 풀 수 없다.

대부분 이러한 구간 값을 탐색하는 문제는 세그먼트 트리나 인덱스드 트리로 풀 수 있다. 하지만 인덱스드 트리가 더 넓은 범위를 커버할 수 있다. 따라서 인덱스드 트리로 이 문제를 접근해볼 것이다.

1. 인덱스트 트리의 구성

인덱스드 트리를 어떻게 구성해야할까?

예제를 통해 살펴보자

처음에 다음과 같은 입력이 주어지면

1맛 사탕 2개 3맛 사탕 3개로 시작하게 된다.

그 때 2번째 사탕을 꺼내면 1이 꺼내지게 되고

1맛 사탕 1개 3맛 사탕 3개로 업데이트가 될 것이다.

그 다음 2번째 맛을 꺼내면 1맛 사탕이 1개뿐이므로 3맛 사탕이 꺼내지게 될 것이다.

그럼 이렇게 생각해볼 수 있겠다.

사탕의 맛에 따른 누적 개수는 해당 맛이 나타낼 수 있는 rank번째의 범위를 나타낸다.

그림으로 나타내면 다음과 같다.

즉 인덱스드 트리가 위의 모양처럼 생겼을 때 1 2 까지는 원본 배열의 첫 번째 인덱스가 나오고 3 4 5 까지는 원본 배열의 3번째 인덱스가 나오는 형태를 가지고 있다.

이 문제는 앞서 인덱스드 트리 문제 (구간 합 구하기 , 구간 곱 구하기)와는 다르게 위에서 부터 데이터를 탐색해야하는 문제였다.

위의 트리에서 2번째를 구하려면 어떻게 이동해야할까? 우선 맨 첫 번째 시작점은 항상 1번 인덱스가 될 것이다.

1번 인덱스의 자식 노드 2개를 보자( 2 번 , 3번 인덱스에 해당한다.) 인덱스드 트리는 항상 Tree[i] = 연산(Tree[i*2], Tree[i*2+1]) 형태로 구성되기 때문에 해당 인덱스의 [인덱스 * 2] , [인덱스 * 2 + 1]에 해당하는 노드들은 현재 인덱스의 자식 노드가 된다.

그래서 1번 인덱스의 자식 노드 2개를 살펴보면 2 와 3번 인덱스 노드에 각각 2 와 3이라는 값이 존재한다. 이 경우 2번으로 이동해야한다. 2라는 뜻은 이쪽 방향에 2개의 사탕이 존재한다는 뜻이기 때문이다.

여기서 한 가지 조건을 알 수 있다.

if(Tree[현재 인덱스 * 2] >= 구하고자 하는 rank) {
	i *= 2; // 왼쪽 자식노드로 이동
} else {
	i = i * 2 + 1;//오른쪽 자식 노드로 이동
}

그리고 한가지 조건이 더 필요하다. 만약 오른쪽 자식으로 이동했을 때는 구하고자 하는 rank의 값이 변해야한다.

만약 4번째 사탕을 구한다고 해보자

그러면 1부터 시작해서 왼쪽 자식 노드의 값보다 구하고자 하는 rank의 값이 크니 오른쪽으로 이동한다. 그럼 여기서도 4번째를 구해야할까? 아니다. 왼쪽 자식 노드의 값이 2이므로 우리는 오른쪽 노드 기준 4-2 = 2 (오른쪽 자식노드로 이동한 후에는 원래 구하고자하는 rank에 오른쪽 자식 노드의 값만큼을 뺴준 순서의 사탕을 구해줘야한다.) 번째 사탕을 찾아야한다. 따라서 재귀함수로 작성한 코드는 다음과 같다.

int query(int rank, int i) {
	if(i >= OFFSET){
    	//offset은 원본 배열의 시작 점
        update(i-OFFSET +1, -1); // 하나씩 빼준다.
        return i - OFFSET + 1; // 사탕의 맛은 1부터 시작이나 원본배열은 OFFSET(짝수)부터 시작하므로
    }
    if(Tree[i * 2] >= rank) {
    	return query(rank, i*2); // 만약 왼쪽 자식 노드의 값이 구하고자 하는 값보다 크거나 같다면
    }else {
    	return query(rank - Tree[i*2] , i*2+1); // 왼쪽 자식 노드의 값이 구하고자 하는 값보다 부족하다면 오른쪽 자식 노드로 이동 
    }
}

또한 update함수는 어떻게 구현해야할까? 인덱스드 트리의 update함수는 거의 bottom-up형태로 구현되어야한다.

따라서 원본배열 위치에 해당하는 인덱스 부분을 업데이트 해주고 해당 인덱스를 인덱스 /= 2해주면서 업데이트해가면 될 거 같다. 따라서 작성된 코드는 다음과 같다.

void update (int id, int cnt) {
	id += OFFSET -1; //원본 배열 위치
    Tree[id] += cnt; // 원본 배열 위치를 우선적으로 업데이트 해준다. 
    while(id > 0) {
    	id /=2; // 그 후 부모 노드로 이동하면서 쭉 업데이트 해준다. 
        Tree[id] = Tree[id * 2] + Tree[id * 2 + 1];
    }
}

4. 코드

#pragma warning (disable :4996)

#include <cstdio>
#define OFFSET 1048576
#define SIZE_TR 1 << 21
using namespace std;

int n; // 사탕상자에 손을 댄 횟수
//comm = 1 사탕 꺼낸다. comm = 2 사탕 cnt만큼 집어넣는다.
int comm; // command
int Tree[SIZE_TR];
void update(int id, int cnt) {
	 id += OFFSET - 1; // 인덱스드 트리의 인덱스는 0부터 시작
	 Tree[id] += cnt;
	 while (id > 0) {
		 id /= 2;
		 Tree[id] = Tree[id * 2] + Tree[id * 2 + 1]; // 부모 노드로 이동하면서 업데이트해준다.
	 }
}

int query(int rank, int id) {
	if (id >= OFFSET) {
		update(id - OFFSET + 1, -1);
		return id - OFFSET + 1;
	}
	if (rank <= Tree[id * 2]) {
		return query(rank, id * 2);
	}
	else {
		return query(rank - Tree[id * 2], id * 2 + 1);
	}
}
void sol() {
	scanf("%d", &n);
	for (int i = 0; i < n; i++) {
		scanf("%d", &comm);
		if (comm == 1) {
			int rank; //
			scanf("%d", &rank); // id에 해당하는 사탕을 꺼낸다. 
			printf("%d\n", query(rank, 1));
		}
		else {
			int id, cnt;
			scanf("%d %d", &id, &cnt);
			update(id, cnt);
		}
	}
}

int main() {
	
	sol();
	return 0;
}