[백준 1321] 군인

문제

www.acmicpc.net/problem/1321

 

알고리즘

세그먼트트리

 

풀이

부대원의 숫자가 계속해서 바뀔 때, 몇번째 병사가 어디 부대에 들어있는지 묻는 문제입니다.

 

누적배열을 만들어 나이브하게 해결해봅시다. 부대원이 변경될때마다 누적배열을 변경하는데에 최대 $O(N)$의 복잡도를 요구하게됩니다. 업데이트를 더 빠르게 할 수 있는 세그먼트 트리를 요구합니다.

 

세그먼트 트리를 사용하면 업데이트와 쿼리가 $O(logN)$에 해결가능합니다. 남은 문제는 누적배열에서 몇번째에 해당하는$K$가 어디에 속한지 찾는 문제입니다. 누적배열은 오름차순의 형태를 띠고 있고 이는 이분탐색을 통해 접근할 수 있다는 의미입니다.

 

펜윅트리에서 쿼리는 누적합을 리턴하므로 펜윅트리 + 이분탐색을 통해 문제를 해결할 수 있습니다.

 

코드

#include <bits/stdc++.h>
#define rep(i,n) for(int i=0;i<n;++i)
#define REP(i,n) for(int i=1;i<=n;++i)
#define FAST cin.tie(NULL);cout.tie(NULL); ios::sync_with_stdio(false)
using namespace std;

const int MAXS = 500'000;


int N,x,Q,a,b,c;
int fenwick[MAXS + 1];

void update(int idx, int val) {
    while (idx <= N) {
        fenwick[idx] += val;
        idx += idx & -idx;
    }
}

int query(int idx) {
    int ret = 0;
    while (idx) {
        ret += fenwick[idx];
        idx -= idx & -idx;
    }
    return ret;
}

int main() {
    FAST;
    cin >> N;
    REP(i, N) {
        cin >> x;
        update(i, x);
    }
    cin >> Q;
    while (Q--) {
        cin >> a;
        if (a == 1) {
            cin >> b >> c;
            update(b, c);
        }
        else {
            cin >> b;
            int lo = 1;
            int hi = N;
            int best;
            while (lo <= hi) {
                int mid = (lo + hi) / 2;
                if (query(mid) >= b) {
                    best = mid;
                    hi = mid - 1;
                }
                else lo = mid + 1;
            }
            cout << best << '\n';
        }
    }

    return 0;
}