[C++ 백준 풀이][Platinum V] 2243번 : 사탕상자

이 포스트에서는 프로그래밍 문제 사이트 백준 Online Judge의 2243번 : '사탕상자' 문제의 풀이 코드와 해설을 다루고 있습니다.

 

문제의 난이도는 Solved.ac 기준 Platinum V에 해당하며, 세그먼트 트리의 응용 개념에 대해 다룰 것입니다.

 

최근 세그먼트 트리 응용과 관련된 문제들을 많이 풀이하고 있는데, 이전 포스트까지는 문제 풀이의 원리에 대해 자세하게 다뤘기 때문에 여기서는 빠르게 풀이만 하고 넘어가도록 하겠습니다.

 

 

2243번 : 사탕상자

 

중복 가능한 번호가 있는 사탕들을 상자에 넣고, 작은 순서로 특정 번째에 있는 사탕을 꺼냈을 때 그 번호를 출력하는 문제입니다.

20억개에 가까운 값들의 수정과 추출이 쿼리로써 반복되는 문제이므로, 이 데이터들을 주어진 시간 안에 처리하기 위해 세그먼트 트리를 사용해야 합니다.

 

Leaf Node가 해당 번호를 가지고 있는 사탕의 갯수를 저장하도록 하고 그 위로 상위 트리를 구성하게 만들어주면 됩니다.

그리고 1번 쿼리에서 왼쪽부터 B번째 사탕까지 합을 구했을 때 B번째 사탕의 번호를 구해야 하는데, 이것은 왼쪽 트리부터 합을 구하면서 왼쪽 트리의 값이 B보다 크거나 같다면 왼쪽 트리를 타고, B보다 작다면 오른쪽 트리를 타서 구간 합을 증가시켜주면 됩니다.

합이 B가 되도록 그 경계선을 탐색하면서 내려가는 방법을 사용하는 것입니다.

 

그럼 문제를 풀이해보겠습니다.

 

 

↓ 위의 이미지와 동일한 풀이 코드는 아래의 접은 글에 정리되어 있습니다.

#include <cstdio>
#include <vector>
#define MAX 1000000
using namespace std;

vector<int> Tree;

void updateTree(int Node, int Begin, int End, int Index, int Diff) {
    if(Index < Begin || Index > End) return;
    Tree[Node] += Diff;
    if(Begin < End) {
        int Mid = (Begin + End)/2;
        updateTree(Node*2, Begin, Mid, Index, Diff);
        updateTree(Node*2+1, Mid+1, End, Index, Diff);
    }
}

int findNum(int Node, int Begin, int End, int Ranking) {
    if(Begin == End) return Begin;
    int Mid = (Begin + End)/2;
    if(Ranking <= Tree[Node*2]) return findNum(Node*2, Begin, Mid, Ranking);
    else return findNum(Node*2+1, Mid+1, End, Ranking-Tree[Node*2]);
}

int main() {
    int N, Q, Ranking, Index, Diff;
    scanf("%d", &N);
    Tree.resize(4*MAX+1);
    for(int i=0; i<N; i++) {
        scanf("%d", &Q);
        if(Q == 1) {
            scanf("%d", &Ranking);
            Index = findNum(1, 1, MAX, Ranking);
            printf("%d\n", Index);
            updateTree(1, 1, MAX, Index, -1);
        }
        else {
            scanf("%d %d", &Index, &Diff);
            updateTree(1, 1, MAX, Index, Diff);
        }
    }
}

 

우선 이 문제에서 주의해야 할 점은 N은 사탕의 최대 번호가 아닌 쿼리의 수입니다.

그렇기 때문에 트리의 크기를 그냥 사탕의 최댓값인 100만*4 이상으로 잡아주어야 합니다.

그리고 각 함수에서 값을 대입할 때도 End의 Index를 MAX로 대입해주어야 합니다.

 

나머지는 대부분 일반적인 세그먼트 트리의 구현과 비슷합니다.

우선 updateTree 함수에서는 특정 Index 구간의 사탕의 개수의 총합을 업데이트 해줍니다.

쿼리 1을 처리할 경우 사탕의 번호를 구하고 그 사탕을 빼게 되므로 해당 Index를 포함하는 모든 노드에 대해 -1개를 더해주면 됩니다.

findNum 함수는 Ranking번째 사탕의 번호를 리턴해주는 함수입니다.

1번 사탕에서부터 (트리의 왼쪽에서부터) Ranking 개수를 세었을 때 나오는 번호를 리턴하면 되므로, 트리를 이분하여 왼쪽 트리의 값이 Ranking보다 크거나 같다면 왼쪽 트리를 타고, 그렇지 않을 경우에는 Ranking에서 왼쪽 트리의 값을 빼고 오른쪽 트리에서 나머지 수를 더 세어주면 됩니다.

 

 

제출해보면 최하위 노드 수가 최대 100만임에도 불구하고 빠른 시간에 프로그램이 작동함을 확인할 수 있습니다.