백준[2243] 사탕상자

문제 설명

문제 링크

http://icpc.me/2243

2243번: 사탕상자

풀이

백준 12899 데이터 구조  이 문제와 거의 동일한 문제다.

세그먼트 트리를 이용하여 k번째 수를 빠르게 구하는 문제다.

업데이트는 일반 세그먼트 트리 업데이트와 비슷하게 하고, 쿼리를 약간 다르게 구현한다.

k번째 수를 구할 때, 왼쪽 자식 노드의 값은 L이라 한다면, k < L이라면 k번째 수는 왼쪽 자식에 있는 것이고, k > L이라면 오른쪽 자식 중 k - L번째 수를 구하는 것이다.

코드

#include <iostream>
#define mid ((s + e) >> 1)

using namespace std;
const int N = 1e6 + 6;

int tree[4 * N];

void update(int node, int s, int e, int k, int c) {
    if (e < k || k < s) return;
    tree[node] += c;
    if (s == e) return;
    update(2 * node, s, mid, k, c);
    update(2 * node + 1, mid + 1, e, k, c);
}

int query(int node, int s, int e, int k) {
    tree[node]--;
    if (s == e) return s;
    if (tree[2 * node] >= k) return query(2 * node, s, mid, k);
    else return query(2 * node + 1, mid + 1, e, k - tree[2 * node]);
}

int main() {
    cin.tie(0) -> sync_with_stdio(0);
    int n; cin >> n;
    while (n--) {
        int a, b; cin >> a >> b;
        if (a == 1) {
            cout << query(1, 0, N, b) << '\n';
        } else {
            int c; cin >> c;
            update(1, 0, N, b, c);
        }
    }
}