[BOJ | 백준] 1321번: 군인

by Kimcoding on 2020-03-26 22:09

in BOJ

문제 설명

BOJ 1321

부대의 감원과 증원으로 군사 재배치가 일어날 때 i번째 군인이 몇 번 부대에 배치 받았는지 구하는 문제입니다.

풀이

세그먼트에서 k번째 요소를 찾는 문제입니다.

유사한 문제로는 BOJ 2243가 있습니다.

세그먼트 트리에 1번~N번 부대에 배치받은 군사 수의 구간합을 저장합니다.

이럴 때 K번째의 군사가 속한 부대 번호를 알아야 합니다.

K번째 군사의 부대 번호를 구하는 쿼리는 루트부터 시작하여

좌, 우 범위의 구간합을 확인한다.
K가 왼쪽에 있는 구간합보다 크다면 오른쪽 범위에 K번째 사람이 존재하므로 오른쪽 범위를 보고 K에서 왼쪽 구간합을 뺍니다.
작거나 같다면 왼쪽 범위에 있으므로 왼쪽 범위를 봅니다.

이 과정을 하다가 리프 노드에 도착하면 리프 노드의 번호가 부대 번호입니다.

소스 코드

#include <bits/stdc++.h>
#define FIO ios::sync_with_stdio(0), cin.tie(0), cout.tie(0)
#define FOR(a,b,c) for(int a = (b); a <= (c); a++)
using namespace std;

struct ST {
	int* tree, leaf;
	ST(int n) {
		leaf = 1 << (int)ceil(log2(n));
		tree = new int[leaf * 2]();
		leaf--;
	}
	void update(int i, int x) {
		i += leaf;
		tree[i] += x;
		while (i >>= 1) tree[i] = tree[i << 1] + tree[i << 1 | 1];
	}
	int query(int x) {
		int i = 1;
		while (i <= leaf) {
			int l = tree[i << 1], r = tree[i << 1 | 1];
			if (l >= x) i <<= 1;
			else i = i << 1 | 1, x -= l;
		}
		return i - leaf;
	}
};

int main() {
	FIO;
	int n; cin >> n;
	ST st(n);
	FOR(i, 1, n) {
		int x; cin >> x;
		st.update(i, x);
	}
	int k; cin >> k;
	while (k--) {
		int t, a; cin >> t >> a;
		if (t == 1) {
			int b; cin >> b;
			st.update(a, b);
		}
		if (t == 2) cout << st.query(a) << '\n';
	}
	return 0;
}