[ BOJ ] 1321 : 군인 ( PLATINUM 5 ) / C

문제

캠프 내내 그랬듯이, 여전히 옆 나라와의 전쟁이 한창이다.

전쟁에는 N개의 부대가 투입되었는데, 전쟁이 장기전이 되다 보니 군사의 적절한 배치를 위해 각 부대에 군인이 늘어나기도 하고 줄어들기도 하고 있다.

행정의 편의를 위해 각 군인들에겐 번호가 붙어 있는데, 군인들은 1번 부대부터 군번순서대로 차례차례 배치된다. 예를 들어 1번 부대에 4명, 2번 부대에 3명, 3번 부대에 7명의 군인이 있다면 군번이 6번인 군인은 2번 부대에 배치되게 된다.

문제는 어떤 부대의 인원이 늘어나거나 줄어들었을 때 i번 군인이 어디에 배치되는지 인데, 이럴 때에는 군인도 군번도 처음부터 다시 배치하게 된다. 위의 예에서와 같이 1번 부대에 4명, 2번 부대에 3명, 3번 부대에 7명의 군인이 있었는데, 1번 부대에서 3명의 감원이 일어난다면, 6번 군인은 3번 부대에 재배치 받게 된다.

전쟁 때는 부대의 감원과 증원이 많아 군사 재배치도 자주 일어나게 되는데, 이렇게 자주 배치가 바뀌자 군인들은 자기가 도대체 어떤 부대에 속하는 지 헷갈리게 되었다. 다행히도 바뀐 군번은 다들 정확하게 숙지하고 있다.

부대의 개수 N과 각 부대에 속해 있는 군인의 수가 N개 주어질 때, 부대의 감원과 증원을 한 후, 혹은 그 중에 군번 i번의 군인이 몇 번 부대에 속하는 지를 물어봤을 때, 그 질문에 대답을 해 줄 수 있는 프로그램을 작성하시오.

i번 부대에 증원이나 감원을 할 때엔 "1 i a"의 형태로 명령이 주어지고, 이는 i번 부대에 a명을 더한다는 뜻이다. 감원을 할 때엔 a가 0보다 작은 수로 주어진다. 감원을 해서 부대의 인원수가 0보다 작아지는 입력은 들어오지 않는다. a는 절댓값이 1보다 크거나 같고, 3,000보다 작거나 같은 정수이다.

군번 i번의 군인이 어떤 부대에 배치 받았는지 알고 싶을 때는 "2 i"의 형태로 명령이 주어지고, 이런 명령을 받았을 때는 i번 군인이 몇 번 부대에 배치 받았는지를 출력해야 한다. i는 전체 군인 수보다 작거나 같은 자연수이다.

입력

첫째 줄에 부대의 개수 N(1 ≤ N ≤ 500,000)이 주어지고, 이어서 각 부대의 군사 수를 나타내는 정수가 N개 주어진다. 각 부대의 군사 수는 1000보다 작거나 같은 자연수이다. 그 다음 줄에 명령의 개수 M(1 ≤ M ≤ 10,000)개가 주어지고, 이어서 M줄에 걸쳐 명령이 주어진다.

출력

질문한 군인이 몇 번 부대에 배치 받았는지를 한 줄에 하나씩 출력한다.

풀이 과정

처음으로 세그먼트 트리에 이분 탐색을 적용했다. 구간 합 트리를 만들고, 왼쪽 노드들의 합 + 현재 노드를 기준으로 value ( 질문한 군인 ) 보다 작으면 end + 1 ( 최소한 이 부대 뒤에 배치되어 있을 것. ), 같으면 end ( 정확히 이 부대에 배치됨. ), 크면 한 단계 더 아래로 탐색 ( 최대한 이 부대 앞에 배치되어 있을 것. ) 하는 이분 탐색을 적용했다.

#include <stdio.h>
#define SIZE 500001
typedef long long LL;

int arr[SIZE];
LL tree[SIZE * 4];

void init(int start, int end, int idx) {
    if (start == end) {
        tree[idx] = arr[start];
        return;
    }

    int mid = (start + end) / 2;
    init(start, mid, idx * 2);
    init(mid + 1, end, idx * 2 + 1);
    tree[idx] = tree[idx * 2] + tree[idx * 2 + 1];
}

void update(int start, int end, int idx, int what, int value) {
    if (what < start || end < what) return;
    if (start == end) {
        tree[idx] += value;
        return;
    }

    int mid = (start + end) / 2;
    update(start, mid, idx * 2, what, value);
    update(mid + 1, end, idx * 2 + 1, what, value);
    tree[idx] = tree[idx * 2] + tree[idx * 2 + 1];
}

LL query(int start, int end, int idx, int value, int leftNode) {
    if (start == end) {
        int result;

        if (leftNode + tree[idx] < value) result = end + 1;
        else if (leftNode + tree[idx] == value) result = end;
        else if (leftNode + tree[idx] > value) {
            if (start == 0) result = 0;
            else result = -1;
        }
        return result;
    }

    int result;
    if (leftNode + tree[idx] < value) result = end + 1;
    if (leftNode + tree[idx] == value) result = end;
    if (leftNode + tree[idx] > value) {
        LL q1 = query(start, (start + end) / 2, idx * 2, value, leftNode);
        if (q1 == -1) return -1;

        LL q2 = query((start + end) / 2 + 1, end, idx * 2 + 1, value, leftNode + tree[idx * 2]);
        if (q1 == -1 && q2 == -1) result = -1;
        else result = (q1 > q2) ? q1 : q2;
    }

    return result;
}

int main(void) {
    int n;
    scanf("%d", &n);
    for (int i = 0; i < n; i++) scanf("%d", &arr[i]);

    init(0, n - 1, 1);

    int m;
    scanf("%d", &m);
    for (int i = 0; i < m; i++) {
        int mode, n1;
        scanf("%d %d", &mode, &n1);
        if (mode == 1) {
            int a;
            scanf("%d", &a);
            update(0, n - 1, 1, n1 - 1, a);
        }
        else {
            printf("%lld\n", query(0, n - 1, 1, n1, 0) + 1);
        }
    }
    return 0;
}