사탕상자

백준 2243번: 사탕상자

아이디어

각 구간에서의 사탕 개수 합을 저장한다.
rank번째로 맛있는 사탕을 사탕을 구할때는 세그먼트 트리에서 binary search를 하면 된다!
루트부터 시작해서, left node에 저장된 합이 rank보다 크거나 같은 경우는 왼쪽으로, rank가 더 큰 경우는 오른쪽으로 가면 된다. 이때 다음 서브트리에서도 똑같은 규칙을 적용하기 위해 rank에서 left node에 저장된 합만큼 빼준다.

코드

#include <bits/stdc++.h>

using namespace std;

int N;
const int MAX = 1000000;
vector<int> tree(4*MAX, 0);

void update(int start, int end, int idx, int diff, int node) {
    if (idx < start || end < idx) return;
    tree[node] += diff;
    if (start == end) return;
    update(start, (start+end)/2, idx, diff, node*2);
    update((start+end)/2+1 ,end, idx, diff, node*2+1);
}

int query(int start, int end, int rank, int node) {
    if (start == end) return start;
    if (tree[node*2] < rank) {
        return query((start+end)/2+1, end, rank-tree[node*2], node*2+1);
    }
    else {
        return query(start, (start+end)/2, rank, node*2);
    }
}

int main() {
    ios_base::sync_with_stdio(0);
    cin.tie(0);
    cin >> N;
    for (int i = 0; i < N; i++) {
        int a, b, c;
        cin >> a;
        if (a == 1) {
            cin >> b;
            int candy = query(0, MAX-1, b, 1);
            cout << candy+1 << '\n';
            update(0, MAX-1, candy, -1, 1);
        }
        else if (a == 2) {
            cin >> b >> c;
            update(0, MAX-1, b-1, c, 1);
        }
    }
    return 0;
}

여담

인덱스 살짝 잘못쓰고, (start+end)/2+1 해야하는데 +1 안해서 또 헤매고,,
헷갈리지 않게 앞으로 start는 무조건 0부터 시작하도록 짜야겠다.