BOJ - 군인(1321)

문제(링크)



구간합 세그먼트 트리를 구현하고, 트리를 update 해가며 풀이합니다.

2번 질의에 대해 루트부터 왼쪽 오른쪽 구간합과 군번을 비교해가며 내려갑니다.

1번부터 봐야하므로 왼쪽 먼저 탐색하며, 왼쪽 구간합에 해당하지 않는다면 오른쪽 구간합을 볼 때 왼쪽 구간합의 값을 빼줍니다.

#include<iostream>
#include<vector>
#include<math.h>
#include<algorithm>
#include<stdlib.h>
#include<memory.h>
using namespace std;

int N, M;
int* tree;

void update(int now, int l, int r, int idx, int diff) {
 if (idx > r || idx < l) return;
 tree[now] += diff;
 if (r == l) return;
 
 update(now * 2, l, (l + r) / 2, idx, diff);
 update(now * 2 + 1, (l + r) / 2 + 1, r, idx, diff);
}

int searchNum(int now, int l, int r, int v) {
 if (r == l) {
  return l;
 }
 if (tree[now * 2] >= v) {
  return searchNum(now * 2, l, (l + r) / 2, v);
 }
 else {
  //오른쪽의 구간합은 왼쪽값이 들어있지 않으므로 (v - 왼쪽 구간합)
  return searchNum(now * 2 + 1, (l + r) / 2 + 1, r, v - tree[now * 2]);
 }
}

int main(void) {
 ios_base::sync_with_stdio(false);
 cin.tie(NULL);

 cin >> N;

 double H = ceil(log2(N));
 tree = (int*)malloc(sizeof(int) * (int)pow(2, H + 1));
 memset(tree, 0, sizeof(int) * (int)pow(2, H + 1));

 for (int i = 1; i <= N; i++)
 {
  int n;
  cin >> n;
  update(1, 1, N, i, n);
 }
 cin >> M;
 for (int i = 0; i < M; i++)
 {
  int q, a, b;
  cin >> q;

  if (q == 1) {
   cin >> a >> b;

   update(1, 1, N, a, b);
  }
  else {
   cin >> a;

   cout << searchNum(1, 1, N, a) << "\n";
  }
 }
}