[BOJ] 12899 : 데이터 구조 (C++)

📖 백준 12899번 : https://www.acmicpc.net/problem/12899

조건

시간 제한	메모리 제한
2 초	512 MB

문제

자연수를 저장하는 데이터베이스 S에 대해 다음의 쿼리를 처리합시다.

유형 1 : S에 자연수 X를 추가한다.

유형 2 : S에 포함된 숫자 중 X번째로 작은 수를 응답하고 그 수를 삭제한다.

입력

첫째 줄에 사전에 있는 쿼리의 수 N 이 주어집니다. (1 ≤ N ≤ 2,000,000)

둘째 줄부터 N개의 줄에 걸쳐 각 쿼리를 나타내는 2개의 정수 T X가 주어집니다.

T가 1이라면 S에 추가할 X가 주어지는 것입니다. (1 ≤ X ≤ 2,000,000)

T가 2라면 X는 S에서 삭제해야 할 몇 번째로 작은 수인지를 나타냅니다. S에 최소 X개의 원소가 있음이 보장됩니다.

출력

유형 2의 쿼리 개수만큼의 줄에 각 쿼리에 대한 답을 출력합니다.

---

풀이

 몇 번째로 작은 수인지를 판단하는 문제이므로 주어진 값 자체를 트리에 저장할 필요가 없다. 값이 등장하는 빈도수를 배열로 잡아 인덱스를 X값으로 생각해 세그먼트 트리를 구상해준다. query에서는 구간의 정보를 찾는 것이 아니라 트리 내부의 값을 기준으로 찾는다. 만약 2번 쿼리의 x값 트리의 왼쪽 구간의 합보다 작다면 그 구간으로 탐색을 지속하고, 더 크다면 오른쪽으로 탐색해 나간다. 오른쪽 구간으로 탐색을 지속할 때, x에서 왼쪽 구간의 값을 빼줘야 오른쪽 구간에서 올바르게 값을 찾아 갈 수 있다.

코드

#include <iostream>
#include <cmath>
#include <vector>
#include <algorithm>

#define sz 2000001

using namespace std;

int n;
int seg[4 * sz];
int arr[sz];

int query(int x, int node = 1, int s = 0, int e = sz - 1) {
    if (s == e) return s;
    int mid = (s + e) / 2;
    int lsum = seg[2 * node];
    if (x <= lsum) return query(x, 2 * node, s, mid);
    else return query(x - lsum, 2 * node + 1, mid + 1, e);
}

void update(int index, int val, int node = 1, int s = 0, int e = sz - 1) {
    if (index < s || index > e) return;
    if (s == e) {
        arr[s] += val;
        seg[node] = arr[s];
        return;
    }
    int mid = (s + e) / 2;
    update(index, val, 2 * node, s, mid);
    update(index, val, 2 * node + 1, mid + 1, e);
    seg[node] = seg[2 * node] + seg[2 * node + 1];
}

int main() {
    ios::sync_with_stdio(false); cin.tie(nullptr);
    cin >> n;
    while (n--) {
        int t, x;
        cin >> t >> x;
        if (t == 1) update(x,1);
        else {
            x = query(x);
            cout << x << '\n';
            update(x, -1);
        }
    }
}