데이터 구조

문제 링크

풀이

처음부터 $2000000$개의 배열에 추가된 수의 개수를 표시하는 방식으로 세그먼트 트리를 구축해 풀 수 있다. 각각의 쿼리 연산을 아래와 같이 처리할 수 있다.

• 데이터 추가 : 주어진 수가 해당하는 배열 인덱스에 +1 ( 세그먼트 트리 리프 노드는 해당 인덱스에 있는 수의 개수를 의미함 )
• 데이터 삭제 : 자신과 자신보다 작은 수의 개수의 합이 주어진 $X$보다 큰 첫 리프 노드를 찾고 그 리프 노드에 -1

코드

#include <iostream>
#include <algorithm>
#include <vector>
#include <cmath>
using namespace std;
using ll = long long;
int n, m, k;
vector<ll> arr, seg;
// n : tree 기준 노드 번호
// t : 업데이트 대상 인덱스
// s : 배열 구간 시작 인덱스
// e : 배열 구간 종료 종료 인덱스 
void add(int n, int t, int s, int e) {
    if (t < s || t > e) return; // 현재 구간에 추가된 수가 없으면 탐색 종료
    seg[n]++; // 현재 구간이 추가된 수를 포함하므로 개수 + 1
    if (s == e) return; // 리프 노드인 경우 탐색 종료
    int mid = (s + e) / 2;
    add(n * 2, t, s, mid); // 왼자식 탐색
    add(n * 2 + 1, t, mid + 1, e); // 오른자식 탐색
}
// n : tree 기준 노드 번호
// s : 배열 구간 시작 인덱스
// e : 배열 구간 종료 종료 인덱스 
// ps : 현재 배열 구간 이전까지의 합
// p : 목표값 X
int del(int n, int s, int e, int ps, int p) {
    // 현재 구간에 제거 대상 인덱스가 없는 경우 탐색 종료. 
    // 1. ps + seg[n] < p : 앞선 구간과 현재 구간의 모든 수를 합쳐도 제거해야할 인덱스 p에 미치지 못하는 경우
    // 2. ps >= p : 현재 구간의 수를 포함하지 않아도 이미 제거해야할 인덱스 p 이상인 경우
    if (ps + seg[n] < p || ps >= p) return 0; 
    seg[n]--; // 현재 구간이 제거된 수를 포함하므로 개수 - 1
    if (s == e) return s; // 리프 노드인 경우 제거된 수 반환
    int mid = (s + e) / 2;
    int s2 = del(n * 2 + 1, mid + 1, e, ps + seg[n * 2], p); // 오른자식 탐색부터 해야함. seg[n * 2] 가 왼자식 탐색에서 변화될 수 있음.
    int s1 = del(n * 2, s, mid, ps, p); // 왼자식 탐색
    return s1 + s2;
}
int main()
{
    ios::sync_with_stdio(0); cin.tie(0); cout.tie(0);
    n = 2000000; arr = vector<ll>(n + 1); seg = vector<ll>((int)pow(2, (int)ceil(log2(n)) + 1));        
    cin >> m;
    for (int i = 1; i <= m; i++) {
        int q, x; cin >> q >> x;
        if (q == 1) {            
            add(1, x, 1, n);
        }
        if (q == 2) {
            cout << del(1, 1, n, 0, x) << "\n";
        }
    }
    return 0;
}