[백준 2243] 사탕상자

 

 세그먼트 트리 문제다. 수정이 빈번하게 일어나는 구간 쿼리 문제는 대부분 세그먼트 트리로 해결할 수 있는 것 같다. 구간 값 update 함수는 일반적인 세그먼트 트리와 같고 query 함수를 조금 신경 쓸 필요가 있다.

 문제에서 요구하는 것은 특정 구간 합 그 자체가 아니라, 꺼내야 하는 사탕의 맛 값이다. 따라서 왼쪽 자식 노드 값, 그러니까 현재 구간 중에서 앞쪽 절반 구간에 존재하는 사탕 수 합이 쿼리하는 값 이상이라면, 해당 구간에 대해 쿼리를 이어나가면 된다. 반대로 앞쪽 절반 구간에 존재하는 사탕 수가 불충분하다면, 당연히 오른쪽 자식을 이용해야 한다. 이때, 주의해야 할 것은 쿼리하는 값을 그대로 사용하면 안 된다는 것이다. 

 전체에서 b번째 맛을 꺼내야 하므로, 앞쪽 구간의 사탕 수만큼을 제외하고 쿼리해야 한다. 만약 필요한 사탕이 14번째 사탕이고 앞쪽 구간에 사탕이 6개가 있다고 치면, 뒤쪽 구간의 8번째 사탕을 구해야 하는 것이다. 쿼리 값을 그대로 사용하게 되면, 전체에서 20번째 사탕을 구하게 될 것이다.

 

 이를 유념하고 작성하면, 큰 틀은 세그먼트 트리의 구현과 같으며 오히려 트리가 비어있는 상태로 시작하므로 초기화 과정에선 공간 할당만 하면 된다. 나는 어차피 query 함수가 사탕을 빼낼 때만 호출될 것이므로, 최적화를 위해 query 함수 내부에서 세그먼트 트리 값들을 -1 하면서 업데이트 해주도록 했다. 쿼리한 값을 가지고 업데이트를 재차 진행하도록 했을 때보다 대략 20ms 가량 더 빠른 결과를 확인할 수 있었다.

 

#include <iostream>

using namespace std;

const int MAX_FAV = 1000000;
long long segTree[MAX_FAV * 4] = {0,};
int n;

int query(int l, int r, int ind, int cnt) {
    segTree[ind]--;
    if(l==r) return l;
    int mid = (l+r)/2;
    if(segTree[ind*2] >= cnt) {
        return query(l, mid, ind*2, cnt);
    } else {
        return query(mid+1, r, ind*2+1, cnt-segTree[ind*2]);
    }
}

void updateTree(int l, int r, int ind, int target, long long delta) {
    if(target<l || target>r) return;
    segTree[ind] += delta;
    if(l==r) {
        return;
    } else {
        int mid = (l+r)/2;
        updateTree(l, mid, ind*2, target, delta);
        updateTree(mid+1, r, ind*2+1, target, delta);
    }
}

int main() {
    ios_base::sync_with_stdio(false);
    cin.tie(nullptr);
    cout.tie(nullptr);

    cin >> n;
    int a, b, c;    

    while(n--) {
        cin >> a >> b;
        if(a==1) {
            int ans = query(1, MAX_FAV, 1, b);
            cout << ans << "\n";
            // updateTree(1, MAX_FAV, 1, ans, -1);
        } else {
            cin >> c;
            updateTree(1, MAX_FAV, 1, b, c);
        }
    }
}