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

https://www.acmicpc.net/problem/2243

 

맛 별로 사탕이 몇 개인지를 세그먼트 트리로 관리해주고, 몇번째 사탕이 무슨 맛인지는 이분 탐색을 통해서 찾았다.

즉, 이 풀이의 시간 복잡도는 로그 제곱이다.

 

참고로, 세그먼트 트리 내에서 왼쪽 구간 합이 b보다 큰지 작은지로 나눠서 처리하도록 하면 로그의 시간 복잡도로 문제를 해결할 수 있다. 이런 유형을 세그먼트 트리에서 k번째 수 찾기라고 한다.

 

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

int n, m;

struct SegTree {
    vector<long long> tree;

    long long update(int idx, int l, int r, int target, long long val) {
        if (target < l || target > r) return tree[idx];
        if (l == r) return tree[idx] = val;
        int m = (l + r) / 2;
        return tree[idx] = update(idx * 2, l, m, target, val) + update(idx * 2 + 1, m + 1, r, target, val);
    }
    
    long long query(int idx, int l, int r, int wl, int wr) {
        if (wr < l || wl > r) return 0;
        if (wl <= l && wr >= r) return tree[idx];
        int m = (l + r) / 2;
        return query(idx * 2, l, m, wl, wr) + query(idx * 2 + 1, m + 1, r, wl, wr);
    }
    
    long long update(int target, long long val) {
        return update(1, 0, n, target, val);
    }
    
    long long query(int wl, int wr) {
        return query(1, 0, n, wl, wr);
    }
};

void solve() {
    SegTree seg;

    n = 1e6 + 1;
    cin >> m;
    seg.tree.assign(4 * n + 1, 0);

    while (m--) {
        int a, b, c;
        cin >> a;
        
        if (a == 1) {
            cin >> b;
            int l = 0;
            int r = n;
            while (l < r) {
                int m = (l + r) / 2;
                if (seg.query(0, m) < b) l = m + 1;
                else r = m;
            }
            cout << l << '\n';
            seg.update(l, seg.query(l, l) - 1);
        } else if (a == 2) {
            cin >> b >> c;
            seg.update(b, seg.query(b, b) + c);
        }
    }
}

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

    int T = 1;
    for (int i = 0; i < T; i++) {
        solve();
    }
}

(AC)