BOJ)1321 군인

문제:icpc.me/1321

군인들이 번호 순서대로 각 부대에 배치되어 있을 때 i번째 군인이 몇번 부대에 배치되어 있는지 확인하는 문제이다.

이 때 부대별로 군인들의 수가 감원이 일어 날 경우 모든 군사들이 재배치를 받게 되므로 세그먼트 트리를 이용하여 갱신과 쿼리를 O(logN)시간에 해결 해 줘야 한다.

쿼리를 처리하는 방법은 노드에 구간 합을 저장하고 있을 때 내가 부대를 찾아야하는 군번을 가지고 왼쪽 자식 노드와 비교를 하여 왼쪽 자식 노드가 더 큰 값을 가지고 있을 경우 내가 찾고자 하는 부대는 왼쪽에 있는 것이므로 왼쪽 자식을 재귀 호출하고, 아닐 경우 오른쪽 자식을 재귀 호출 하는데 이 때 우리는 왼쪽 자식 수만큼을 제외하고 봐야 하므로 찾아야하는 군번에 왼쪽 자식의 수만큼 빼준 값으로 재귀 호출해준다. 이후 리프에 도달할 경우 해당 노드가 담당하는 구간 번호를 출력해주면 된다.

#include <cstdio>
#include <algorithm>
#define MAX_N 500000
using namespace std;
int n, m, seg[4 * MAX_N], arr[MAX_N], a, b, c;
int update(int pos, int val, int node, int x, int y) {
    if (y < pos || pos < x)
        return seg[node];
    if (x == y)
        return seg[node] = val;
    int mid = (x + y) >> 1;
    return seg[node] = update(pos, val, node * 2, x, mid) + update(pos, val, node * 2 + 1, mid + 1, y);
}
int query(int val, int node, int x, int y) {
    int mid = (x + y) >> 1;
    if (x == y)
        return x;
    if (seg[node * 2] >= val)
        return query(val, node * 2, x, mid);
    return query(val - seg[node * 2], node * 2 + 1, mid + 1, y);
}
int main() {
    scanf("%d", &n);
    for (int i = 1; i <= n; i++) {
        scanf("%d", &arr[i]);
        update(i, arr[i], 1, 1, n);
    }
    scanf("%d", &m);
    while (m--) {
        scanf("%d", &a);
        if (a == 1) {
            scanf("%d%d", &b, &c);
            arr[b] += c;
            update(b, arr[b], 1, 1, n);
        }
        else {
            scanf("%d", &b);
            printf("%d\n", query(b, 1, 1, n));
        }
    }
    return 0;
}