백준 - 사탕상자 (2243) [C++]

이 문제는 세그먼트 트리를 이용하여 풀었다.

 

이 문제의 핵심은 segment의 리프 노드를 1 ~ 1,000,000의 맛 사탕의 개수를 나타내는 것으로 생각할 수 있는지이다. 이걸 아직 생각 못했다.

또한 두 번째는 query를 호출하고, cntInSection이 segment[node] 가 아닌 segment[node * 2]를 저장해야한다는 것이다.

segment[node]로 해놓을 경우 실제로 rank가 제대로 계산되기 전에 start == end에 도착해버리고, 한 단계 낮은 값이 나올 가능성이 있다는 문제가 있다. 따라서 먼저 왼쪽 노드를 확인함으로써 이러한 문제를 해결할 수 있다.

 

한 가지 유의할 점은 A == 1일 때 query()를 호출한 값으로 update를 실행해야한다는 점이다. 나는 query()와 update() 모두에 B를 넣었다가 원인도 모르고 계속 틀렸다.

#include <iostream>
#include <algorithm>
using namespace std;

#define MAX 1000000

typedef long long ll;

ll segment[MAX * 4 + 1];
int n;

int query(int node, int start, int end, ll rank) {
    if (start == end) {
        return start;
    }

    int mid = (start + end) / 2;
    ll cntInSection = segment[node * 2];
    
    if (cntInSection >= rank) {
        return query(node * 2, start, mid, rank);
    } else {
        return query(node * 2 + 1, mid + 1, end, rank - cntInSection);
    }
}

void update(int node, int start, int end, int index, int diff) {
    if (index < start || end < index) {
        return;
    }

    segment[node] += diff;

    int mid = (start + end) / 2;
    if (start != end) {
        update(node * 2, start, mid, index, diff);
        update(node * 2 + 1, mid + 1, end, index, diff);
    }
}

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

    cin >> n;
    for (int i = 0; i < n; i++) {
        int A, B, C;

        cin >> A;

        if (A == 1) {   // pop
            cin >> B;

            int ans = query(1, 1, 1000000, B);
            cout << ans << '\n';
            update(1, 1, 1000000, ans, -1);
        } else {
            cin >> B >> C;

            update(1, 1, 1000000, B, C);
        }
    }

    return 0;
}