[baekjoon] 2243 사탕 상자 (cpp)

2243 사탕 상자 link

문제

---

풀이

---

문제를 정리하면

1. 배열에 사탕이 들어있고, 사탕의 맛은 숫자로 표현된다. (최대 2,000,000,000개)
2. A가 1인 query가 들어오면 배열에서 B번째로 맛있는 (B번째로 작은) 사탕을 꺼내준다. (출력한다)

사탕의 맛은 1 ~ 1,000,000 이므로 백만개 길이의 배열을 두고 각 Index에 그 맛을 가진 사탕을 저장하는 방법을 생각해볼 수 있다.
그런 후에 사탕을 꺼내려면 앞에서부터 B번째 사탕의 맛을 찾아 떠나야한다.
그러면 query가 최대 100,000, 배열 최대 길이가 1,000,000이니까... 무조건 시간초과고 더 효과적인 방법을 찾아야한다.

조금만 생각을 바꿔보자.
앞에서부터 B번째 사탕을 찾는다 = 낮은 숫자부터 사탕이 몇개씩 있나 누적합을 구해야한다.
누적 하면 역시 세그먼트 트리

이 아이디어를 가지고 풀이를 정리해보면

1. 세그먼트 트리에 사탕 맛 구간별 사탕의 개수를 구해둔다.
2. 1 ~ x 구간의 사탕 개수가 B가 되는 x 가 원하는 사탕 맛일 것이다.

근데 2번에서 x를 구할 때 1부터 순차적으로 계산하면 linear하기 때문에 세그먼트 트리를 만든 보람이 없다.
따라서 1은 고정해두고 이분탐색으로 x를 구하면 시간을 단축할 수 있다.

Code

#include <iostream>
#define MAXN 1000000

int tree[1<<21];
int n;

int query(int node,  int s, int e, int l, int r) {
  if (r < s || e < l) 
    return 0;

  if (l <= s && e <= r) 
    return tree[node];

  int mid = (s + e) / 2;
  return query(node * 2, s, mid, l, r) + query(node * 2 + 1, mid + 1, e, l, r);

}

void update(int node, int s, int e, int idx, int val) {
  if (idx < s || e < idx)
    return;

  tree[node] += val;
  if (s == e) return;

  int mid = (s + e) / 2;
  update(node * 2, s, mid, idx, val);
  update(node * 2 + 1, mid + 1, e, idx, val);
}

int main() {
  scanf("%d", &n);
  for (int i = 0; i < n;i++) {
    int a, b, c;
    scanf("%d %d", &a, &b);
    if (a == 1) {
      int s = 1, e = MAXN;
      while (e - s > 0) {
        int mid = (s + e) / 2;
        int q = query(1, 1, MAXN, 1, mid);
        if (q < b)
          s = mid + 1;
        else 
          e = mid;
      }
      printf("%d\n", e);
      update(1, 1, MAXN, e, -1);
    } else if (a == 2) {
      scanf("%d", &c);
      update(1, 1, MAXN, b, c);
    }
  }
}

사실 로직은 완전 동일하나 파이썬으로 하니 자꾸 시간초과가 나서... 그냥 cpp로 하니 통과했다.
편리함에는 대가가 있는 법...ㅠㅠ