[Segtree]BOJ 2243_사탕 상자

문제 링크 : www.acmicpc.net/problem/2243

2243번: 사탕상자

문제의 상황에서 맛있는 정도가 1인 사탕이 3개, 2인 사탕이 4개, 5인 사탕이 10개 있다고 가정하자.

이 때, 6번째 사탕과 12번째 사탕의 맛있는 정도는 얼마일까?

 

k번째 사탕의 맛있는 정도를 알기 위해서는 각 사탕들 갯수에 대한 누적합이 필요하다.

사탕이 맛있는 정도를 d라 하자. d = 1 인 사탕은 3개, d = 2인 사탕은 4개이다.

d = 2인 사탕까지의 누적합은 7이고, 이는 6보다 큰 최소의 누적합이므로

우리는 6번째 사탕의 d가 2임을 알 수 있다.

마찬가지로, 12번째 사탕의 d가 3임도 알 수 있다.

 

누적합을 배열로 저장해 갱신하며 사용하는 경우는 시간복잡도 면에서 매우 불리하다.

사탕이 맛있는 정도인 d가 100만까지고, k번째 사탕의 d를 묻는 쿼리는 최대 10만번이다.

배열로 구현할 경우, 최악의 경우 1천억 번의 계산을 하게 되는 것이다.

 

따라서, 세그먼트 트리를 이용한다. 세그먼트 트리는 이 게시글에 정리되어 있다.

세그먼트 트리(Segment Tree)

 

일반적인 세그먼트 트리의 쿼리에서 약간의 변형을 가해주면 쉽게 해결할 수 있다.

결국 k번째 사탕의 d를 알기 위해서는 누적합이 필요하므로,

쿼리에서 트리의 각 노드를 타고 내려갈 때 누적합을 따로 저장해주면 된다.

 

[코드]

#define _CRT_SECURE_NO_WARNINGS
#include <iostream>
#include <vector>
#include <memory.h>
 
#define MAX_TREESIZE 1000000 * 4
#define LEAF 1000000 
 
using namespace std;
int n;
int segtree[MAX_TREESIZE];
 
 
void add_candy(int node, int start, int end,  int target, int val) {
    //segment tree의 target 에 해당하는 리프 노드와 그 상위 노드들에 val만큼 값을 더함
    if (start > target || end < target) return;
 
    segtree[node] += val;
    if (start == end && start == target) return;
 
    int mid = (start + end) / 2;
    add_candy(node * 2, start, mid, target, val);
    add_candy(node * 2 + 1, mid + 1, end, target, val);
}
 
void get_candy(int node, int start, int end, int target, int acc) {
    //누적합을 통해 target 번째의 사탕이 몇 번인지 확인
    if (start == end) {
        //도달한 곳이 leaf 노드라면
        printf("%d\n", start);
        add_candy(1, 1, LEAF, start, -1);
        return;
    }
    int mid = (start + end) / 2;
    if (acc + segtree[node * 2] >= target)
        //한 노드에 대해 왼쪽 자식의 구간합을 확인해, acc + left의 값이 target보다 작다면 left, 아니라면 right 노드로 이동한다.
        get_candy(node * 2, start, mid, target, acc);
    else
        get_candy(node * 2 + 1, mid + 1, end, target, acc + segtree[node * 2]);
    
}
 
void init() {
    int cmd, a1, a2;
    cin >> n;
    for (int i = 0; i < n; i++) {
        cin >> cmd;
        if (cmd == 1) {
            cin >> a1;
            get_candy(1, 1, LEAF, a1, 0);
        }
        else {
            cin >> a1 >> a2;
            add_candy(1, 1, LEAF, a1, a2);
        }
    }
 
}
 
int main() {
    ios::sync_with_stdio(0);
    cin.tie(0);
    cout.tie(NULL);
 
    init();
 
    return 0;
}