[백준] 2243번 사탕상자 C++ 코드

코드

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

using namespace std;
typedef long long ll;
ll MAX_N = 1000003;
ll n, a;
vector<ll> tree;

void update(vector<ll> &tree, ll idx, ll diff) {
    while(idx <= MAX_N) {
        tree[idx] += diff;
        idx += (idx & -idx);
    }
}

ll sum(vector<ll> &tree, ll idx) {
    ll ans = 0;
    while(idx > 0) {
        ans += tree[idx];
        idx -= (idx & -idx);
    } 
    return ans;
}
// k번째를 찾는 과정
// 사탕은 인덱스별로 저장되어있고, 개수가 저장되어 있기 때문에 정렬된 상태
// 이분탐색을 통해서 개수가 많으면 줄여서 찾고
// 적으면 s를 올려서 탐색
ll find_kth(ll x) {
    ll s = 1;
    ll e = MAX_N;
    while(s < e) {
        ll mid = (s + e) / 2;
        ll cnt = sum(tree, mid);
        if(cnt >= x) e = mid;
        else s = mid + 1;
    }
    return e;
}

void FASTIO() {
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    cout.tie(NULL);
}

int main() {
    FASTIO();
    cin >> n;
    tree.resize(MAX_N, 0);
    for(int i = 0; i < n; i++) {
        cin >> a;
        if(a == 1) {        
            // 꺼내는 작업
            int b;
            cin >> b;
            int kth = find_kth(b);
            cout << kth << '\n';
            update(tree, kth, -1);
        } else if(a == 2) {
            // 사탕을 넣거나 빼는 작업
            int b,c;
            cin >> b >> c;
            update(tree, b, c);
        }
    }
    return 0;
}

풀이

사탕의 맛은 1부터 1000000까지 구분된다. 맛의 정도에 따라 구분을 해주고 개수를 세어준다. 내부 변화가 많으므로 펜윅트리를 통해 구성했다. k번째 사탕을 찾을 때, 이분탐색을 사용해서 찾을 수 있다. 기본적으로 정렬된 상태이며 개수를 찾아주면 되기 때문에 start와 end를 조정하며 몇 번째 사탕을 구하면 되는지 확인하면 된다. 사탕을 꺼내준 후에는 -1개를 해주어야 하는 걸 잊으면 안된다!