백준 2263 - 사탕상자

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

 

특정 맛을 가진 사탕의 개수를 변경하거나 k번째로 맛있는 사탕을 꺼내야 하는 질의를 최대 10만개 수행해야 한다.

 

세그먼트 트리를 통해 각 맛의 정도에 해당하는 사탕의 개수를 세그먼트 트리로 관리하면서 k번째로 맛있는 사탕(이하 k번째 사탕)을 추적하는 질의를 작성하면 된다.

 

맛의 정도를 index로 하여 배열을 구성하면 update는 구간합 용도로 구현한 세그먼트 트리의 그것과 동일하게 작성이 가능하다.

 

k번째 사탕을 추적하는 질의는 현재 노드에서 왼쪽 구간의 사탕 개수와 오른쪽 구간의 사탕 개수를 비교한다.

 

만약 k보다 왼쪽 구간의 개수 이상이면 왼쪽 노드, 그렇지 않으면 오른쪽 노드로 가지치기 하여 로그 스케일로 줄여나가 원하는 사탕을 찾으면 된다.

 

오른쪽 노드로 가지치기를 할 때는 왼쪽 구간의 사탕을 배제해야 하므로 k에 왼쪽 구간의 사탕 개수를 빼줘야 함을 유의한다.

 

사탕의 맛의 상한이 $10^{6}$임을 주의하여 질의를 작성하면 된다.

 

전체 코드

#include <bits/stdc++.h>

using namespace std;

int n;

inline int mid(int stt, int end) { return stt + (end - stt) / 2; }

void update(vector<int> &tree, int node, int start, int end, int idx, int diff)
{
    if (idx < start || idx > end) return;
    
    tree[node] += diff;
    if (start != end)
    {
        int m = mid(start, end);
        update(tree, node * 2, start, m, idx, diff);
        update(tree, node * 2 + 1, m + 1, end, idx, diff);
    }
}

int query(vector<int> &tree, int node, int start, int end, int kth)
{
    if (start == end) return start;
    int candy_left = tree[node * 2];
    int m = mid(start, end);

    if (candy_left >= kth) return query(tree, node * 2, start, m, kth);
    else return query(tree, node * 2 + 1, m + 1, end, kth - candy_left);
}

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

    cin >> n;

    vector<int> segtree(1e6 * 4 + 4);

    while (n--)
    {
        int qry;
        cin >> qry;

        if (qry == 1)
        {
            int kth;
            cin >> kth;
            int candy = query(segtree, 1, 1, 1e6, kth);
            cout << candy << '\n';
            update(segtree, 1, 1, 1e6, candy, -1);
        }
        else
        {
            int tasty, cnt;
            cin >> tasty >> cnt;
            update(segtree, 1, 1, 1e6, tasty, cnt);
        }
    }
}