[ 백준 2243 ] 사탕상자 (C++)

백준의 사탕상자(2243) 문제이다.

[ 문제 바로가기 ]

[ 문제풀이 ]

입력으로 주어지는 우선순위에 맞게 사탕을 꺼내주기도 하고 넣어주기도 하는 과정을 구현해야 하는 문제이다.

본인은 이 문제를 세그먼트 트리를 이용해서 접근해 보았다. 세그먼트 트리는 '구간에 대한 연산'을 효율적으로 처리할 수 있는 자료구조이다. 이를 적용시키기 위해서 "사탕의 갯수"를 저장하는 세그먼트 트리로 구현해보았다. 즉, 세그먼트 트리의 리프노드에는 사탕의 맛이 i인 사탕의 갯수가 저장되어 있을 것이다.

그럼 본격적인 문제를 풀기전에 세그먼트 트리에 대해 잘 모른다면 아래의 글을 읽고 오도록 하자.

[ 세그먼트 트리 알아보기(Click) ]

위의 링크가 걸려있는 글에는 "구간에 대한 합"을 기준으로 세그먼트 트리가 설명되어 있다. 이제 이를 이 문제에 알맞게 적절히 바꿔보도록 하자.

1. 트리의 생성과정

일반적인 문제들 처럼 배열이 주어지고 구간에 대한 연산을 처리하라는 명령어가 떨어진다면, 해당 배열이 해야 하는 연산 결과값을 세그먼트 트리의 노드에 저장함으로써 세그먼트 트리를 생성하게 된다. 하지만 이 문제 같은 경우 "초기의 사탕상자는 빈 상태로 시작한다" 라고 했으므로 별다르게 세그먼트 트리의 노드 값들을 설정해줄 필요가 없다. 따라서 본인은 세그먼트 트리가 만들어질 크기만 계산해서 공간만 할당해주었다.

2. 사탕을 꺼내는 연산

문제의 입력으로 주어지는 "x번" 째 사탕을 꺼내야 하는 연산을 처리해주어야 한다. 세그먼트 트리에서는 현재 노드에서 왼쪽 자식노드와 오른쪽 자식 노드로 나누어서 탐색을 진행한다. 그런데 이 과정에서 "x번"째 있는 사탕을 꺼내야 한다.

"x번째 사탕"이라는 것은 어떻게 보면 "해당 사탕보다 앞에 x - 1 개의 사탕이 있다" 라고 볼 수 있다.

예를 들어서 { 1, 2, 3, 4, 5 } 이렇게 사탕이 존재하고, 여기서 3번째 사탕을 꺼내주세요 라고 했을 때, 3번째 사탕 앞에는 2개의 사탕이 더 있다는 것이다. 즉 ! 우리는 이 'x - 1개'를 통해서 사탕을 꺼내는 연산을 해 볼 것이다.

세그먼트 트리에서 왼쪽 자식 오른쪽 자식 노드로 나누어서 탐색을 진행할 것인데, 해당 노드에는 사탕의 갯수가 저장되어 있다고 말했다. 그렇다면 현재 노드의 왼쪽 자식에 있는 사탕의 갯수가 'x번' 보다 작다면 어떻게 될까 ??

쉽게 생각해서 "나는 지금 3번째 사탕을 뽑고 싶은데, 내 왼쪽자식에는 사탕이 1개 밖에 없는 상황" 인 것이다. 그럼 이 경우에 왼쪽 자식을 탐색해본다고 해서 3번째 사탕을 뽑을 수 있을까 ?? 절대로 뽑을 수가 없다. 따라서 이런 경우에는 왼쪽 자식을 탐색해볼 필요가 없는 것이다. 그렇다면 오른쪽 자식으로 가야하는데, 오른쪽 자식으로 갈 때는 어떻게 될까 ??

"나는 지금 3번째 사탕을 뽑고 싶은데, 왼쪽 자식에는 사탕이 1개밖에 없어. 그래서 오른쪽 자식으로 가려는데, 오른쪽 자식에서 몇 번째 사탕을 뽑으면 돼?" 라고 묻는 상황이다. 바로 이 경우에는 오른쪽 자식에서 '2번째' 사탕을 뽑으면 된다.

왜냐하면 ! 이미 왼쪽 자식에 '1개'의 사탕이 있기 때문에, 그리고 왼쪽 자식은 오른쪽 자식에 비해서 상대적으로 더 작은 구간을 가지고 있는 노드이기 때문에, 왼쪽 자식에 '1개'가 있으면, 이 '1개'를 포함해서 오른쪽 자식에서는 '2번째'사탕을 뽑아야 하는 것이다. 그래야지 결과적으로 봤을 때 우리가 원하는 '3번째 사탕'을 꺼내게 되는 것이다.

위에서 말한 부분이 Query부분이다. 본인은 다음과 같이 코드로 구현해 보았다.

int Query(int Node, int Start, int End, int Cnt)
{
    if (Start == End) return Start;
 
    int Mid = (Start + End) / 2;
    if (SegmentTree[Node * 2] >= Cnt) return Query(Node * 2, Start, Mid, Cnt);
    return Query(Node * 2 + 1, Mid + 1, End, Cnt - SegmentTree[Node * 2]);
}

line6 ~ 7) 에서 보면 "왼쪽 자식을 탐색하기 위한 조건"이 걸려있다. 바로 내가 찾고자 하는 'x번째'에 해당하는 x 값보다 더 많은 사탕을 가지고 있는 경우에만 왼쪽 자식을 탐색한다. 그게 아니라면 오른쪽 자식을 탐색할 것인데, 'x - 왼쪽자식이 가지고 있는 사탕의 수'를 한 값을 매개변수로 호출하게 된다.

위의 Query를 통해서 'x번째 사탕'을 뽑았을 때, 그 사탕이 몇 번째 맛인지는 구할 수 있다. 그런데 ! 우리는 한 가지 연산을 더 진행해주어야 한다. 바로 Update하는 과정이다. 왜냐하면 사탕 하나를 꺼내 주었기 때문에 결과적으로 본다면 사탕의 갯수는 1만큼 감소하게 된 것이고, 이 1만큼 감소하게 것에 대한 Update를 처리해 주어야 한다.

Update하는 과정은 기존의 세그먼트 트리에서 Update하는 과정과 크게 다른 부분이 없다. 따라서 이 부분에 대한 설명은 위의 세그먼트 트리에 대한 설명 글로 대체하겠다.

[ 소스코드 ]

#include<iostream>
#include<vector>
#include<cmath>
 
#define endl "\n"
#define MAX 1000000
typedef long long ll;
using namespace std;
 
int N;
int Arr[MAX + 5];
vector<ll> SegmentTree;
vector<pair<int, pair<int, int>>> Cmd;
 
void Input()
{
    cin >> N;
    for (int i = 0; i < N; i++)
    {
        int a; cin >> a;
        if (a == 1)
        {
            int b; cin >> b;
            Cmd.push_back(make_pair(a, make_pair(b, -1)));
        }
        else
        {
            int b, c; cin >> b >> c;
            Cmd.push_back(make_pair(a, make_pair(b, c)));
        }
    }
}
 
int Query(int Node, int Start, int End, int Cnt)
{
    if (Start == End) return Start;
 
    int Mid = (Start + End) / 2;
    if (SegmentTree[Node * 2] >= Cnt) return Query(Node * 2, Start, Mid, Cnt);
    return Query(Node * 2 + 1, Mid + 1, End, Cnt - SegmentTree[Node * 2]);
}
 
void Update(int Node, int Start, int End, int Idx, int Diff)
{
    if (Idx < Start || Idx > End) return;
    SegmentTree[Node] = SegmentTree[Node] + Diff;
 
    if (Start != End)
    {
        int Mid = (Start + End) / 2;
        Update(Node * 2, Start, Mid, Idx, Diff);
        Update(Node * 2 + 1, Mid + 1, End, Idx, Diff);
    }
}
 
void Solution()
{
    int Tree_Height = (int)ceil(log2(MAX));
    int Tree_Size = 1 << (Tree_Height + 1);
    SegmentTree.resize(Tree_Size);
    
    for (int i = 0; i < Cmd.size(); i++)
    {
        int Command = Cmd[i].first;
        if (Command == 1)
        {
            int Idx = Cmd[i].second.first;
            int Pos = Query(1, 1, MAX, Idx);
            cout << Pos << endl;
            Arr[Pos]--;
            Update(1, 1, MAX, Pos, -1);
        }
        else
        {
            int Idx = Cmd[i].second.first;
            int Value = Cmd[i].second.second;
            Arr[Idx] = Arr[Idx] + Value;
            Update(1, 1, MAX, Idx, Value);
        }
    }
}
 
void Solve()
{
    Input();
    Solution();
}
 
int main(void)
{
    ios::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
 
    //freopen("Input.txt", "r", stdin);
    Solve();
 
    return 0;
}