[2243번] 사탕 상자

문제 출처 :

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

알고리즘 분석 :

문제 해결에 필요한 사항

1. 세그먼트 트리

이 문제는 펜윅 트리로는 query문을 제작 할 수 없어 세그먼트 트리로 해결하였다.

우선 사탕 종류가 100만가지라는 것을 토대로

리프 노드가 100만개인 세그먼트 트리를 미리 바로 만들어버린다.

그리고 update부분은 일반 세그먼트 트리와 같은 방식으로 update를 해주면 되고

query부분이 색다른 방식이다.

ll query(vector<ll> &tree, int node, int start, int end, int k)
{
    if ((start == end) && ret == 0)
    {
        printf("%d\n", start);
        return start;
    }
 
    // 자신의 왼쪽 자식 노드가 k개 이상 가지고 있다면 왼쪽으로 탐색
    if (ret == 0 && (node*2 <= tree_size && tree[node * 2] >= k))
        return ret = query(tree, node * 2, start, (start + end) / 2, k);
 
    k -= tree[node*2];
 
    // 그 외에는 오른쪽 자식 노드로 탐색
    if(ret == 0 && (node * 2 + 1 <= tree_size && tree[node*2+1] >= k))
        return ret = query(tree, node * 2 + 1, (start + end) / 2 + 1, end, k);
}

if((start == end) && ret == 0)이라는 부분에서 ret는 한번 출력됐으면 더이상 출력을 막기위해 만들어 두었다.

어짜피 사탕의 수는 리프 노드에 모두 있기 때문에 query를 타면 리프 노드로 무조건 향해가야한다.

이때 리프 노드로 가는 조건이 있다.

k번째 사탕을 알기위해선 세그먼트 트리를 조사해 볼 필요가 있다.

우선 arr[] = {1,3,4,2}로 구성 된 세그먼트 트리가 있다고 가정하자

　　　　10

　　4　　　　6

1　　　3　4　　2

이렇게 세그먼트 트리가 형성되면 루트 노드에는 arr의 전체 합인 10이 있고 그 아래는 각각의 누적합들이 존재한다.

(node*2 <= tree_size && tree[node * 2] >= k)에 대해 생각해보자.

일단 node*2가 당연히 세그먼트 트리의 최대 노드 수를 벗어나면 안되는 것이고

tree[node*2] >= k라함은 tree[node*2] 부분에 사탕의 수가 k개 보다 많아야 k번째 수를 찾을 수 있다는 것이다.

만약 저 조건에 해당하지 않는다면

k -= tree[node*2];를 해주는데 이 의미는 다음 설명을 보면 이해하기 쉽다.

---

만약 1번 사탕이 3개, 2번 사탕이 3개 있었다 보자.

1 1 1 2 2 2

이때 3번째 사탕은 1이다.

4번째 사탕은 2임을 알 수 있다.

이를 세그먼트 트리에서 왼쪽 노드 오른쪽 노드로 생각하기 위해서는

1 1 1사탕(왼쪽 노드 수)을 빼고 오른쪽 노드에서 1번째 사탕을 찾으면 된다는 것이다.

이 말이 k = k - tree[node*2]이다. (1(바뀐 k번 째) = 4 (원래 k번째)- 3(왼쪽노드))

---

이렇게 탐색을 하여 query를 조사하면 무조건 적으로 start == end를 들어가게 된다.

소스 코드 : 

#include <iostream>
#include <cstdio>
#include <cmath>
#include <vector>
 
#define MAX_CANDY 1000000
 
using namespace std;
 
typedef long long ll;
 
int k;
int tree_size;
int ret;
 
void update(vector<ll> &tree, int node, int start, int end, int index, ll diff)
{
    if (!(start <= index && index <= end))
        return;
 
    tree[node] += diff;
 
    if (start != end)
    {
        int mid = (start + end) / 2;
        update(tree, node * 2, start, mid, index, diff);
        update(tree, node * 2 + 1, mid + 1, end, index, diff);
    }
}
 
ll query(vector<ll> &tree, int node, int start, int end, int k)
{
    if ((start == end) && ret == 0)
    {
        printf("%d\n", start);
        return start;
    }
 
    // 자신의 왼쪽 자식 노드가 k개 이상 가지고 있다면 왼쪽으로 탐색
    if (ret == 0 && (node*2 <= tree_size && tree[node * 2] >= k))
        return ret = query(tree, node * 2, start, (start + end) / 2, k);
 
    k -= tree[node*2];
 
    // 그 외에는 오른쪽 자식 노드로 탐색
    if(ret == 0 && (node * 2 + 1 <= tree_size && tree[node*2+1] >= k))
        return ret = query(tree, node * 2 + 1, (start + end) / 2 + 1, end, k);
}
 
int main()
{
    int n;
    int h = (int)ceil(log2(MAX_CANDY));
    tree_size = (1 << (1 + h));
 
    vector<ll> tree(tree_size);
 
    scanf("%d", &n);
 
    while (n--)
    {
        int num;
        scanf("%d", &num);
 
        if (num == 1)
        {
            scanf("%d", &k);
 
            int getIndex = query(tree, 1, 0, MAX_CANDY - 1, k);
            ret = 0;
            update(tree, 1, 0, MAX_CANDY - 1, getIndex, -1);
        }
 
        else if (num == 2)
        {
            int index;
            ll val;
            scanf("%d %lld", &index, &val);
 
            update(tree, 1, 0, MAX_CANDY-1, index, val);
        }
    }
 
    return 0;
}
 
//                                                       This source code Copyright belongs to Crocus
//                                                        If you want to see more? click here >>