[알고리즘] 백준 1321번: 군인 (C++)

https://www.acmicpc.net/problem/1321

1321번: 군인

N번 부대에 필요한 사람 수만큼 사람을 1번부터 차례대로 집어넣는다.

어떠한 X번 군번이 주어졌을 때 이 사람은 어느 부대를 들어가야 하는지 알아내야 하는 문제다.

1번 부대부터 필요한 사람 수들을 더하면서 만족하는 범위 안에 들어가면 답을 도출할 수 있다.

하지만, 쿼리는 M개 부대의 수는 N개로 O(N * M)

>> O(N * M), 500,000 * 10,000이 되어 시간 초과가 난다.

 

구간과 관련되었으니까 구간 합 세그먼트 트리를 활용해보자. 매번 쿼리에서 1 ~ n까지 1부터의 쿼리를 차례대로 구하면서

답을 도출할 수 있다.

>> 이럴 경우 시간 복잡도는 O(M * N * logN)

 

오히려 시간이 더 늘었다. 이러면 세그먼트 트리 또한 답이 아닐까? 아니다. 1 ~ n을 선형적으로 다루는 쿼리가 지금 주된 원인이다. 즉, 다루는 범위가 넓기 때문에 특정한 인덱스를 찾기에 시간이 꽤 걸리는 것이다.

>> 방금 한 말에서 이분 탐색을 써야겠다는 생각을 할 수 있다.

>> 시간 복잡도는 O(M * logN * logN)

 

이분 탐색을 먼저 떠올렸다한들 그 구간합을 과정에서도 O(N)이 발생해서 구간 합 세그먼트 트리를 찾았을 거 같다.

 

이분 탐색을 돌려서 나오는 인덱스들이 X 군번보다 작거나 많거나 중 조건을 분기하여 이분 탐색을 돌리면 된다.

매번 1 ~ mid까지 구간 합을 돌려서 주어진 person보다 작다면 low를 높게 주고, 돌린다.

high는 반대의 경우로 마찬가지

이분 탐색을 통해 나온 answer부대에는 person보다 작은 부대이기 때문에 person이 들어가는 부대는 answer + 1부대다. >> answer + 1을 return

int binarySearch(int person) { // Person == 군번
	int low = 1;
	int high = n;
	int answer = 0;
	while (low <= high) {
		int mid = (low + high) / 2;

		if (query(1, n, 1, 1, mid) < person) {
			low = mid + 1;
			answer = mid;
		}
		else {
			high = mid - 1;
		}
	}
	return answer + 1;
}

 

[AC 코드]

#include <iostream>
#define MAX (500000 + 1)
using namespace std;

int arr[MAX];
int sgTree[MAX * 4];
int n;

int sgInit(int start, int end, int node) {
	if (start == end) return sgTree[node] = arr[start];
	int mid = (start + end) >> 1;
	return sgTree[node] = sgInit(start, mid, node * 2) +
		sgInit(mid + 1, end, node * 2 + 1);
}

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

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

int query(int start, int end, int node, int left, int right) {
	if (start > right || end < left) return 0;
	if (left <= start && end <= right) return sgTree[node];
	int mid = (start + end) >> 1;
	return query(start, mid, node * 2, left, right) + 
		query(mid + 1, end, node * 2 + 1, left, right);
}

int binarySearch(int person) {
	int low = 1;
	int high = n;
	int answer = 0;
	while (low <= high) {
		int mid = (low + high) / 2;

		if (query(1, n, 1, 1, mid) < person) {
			low = mid + 1;
			answer = mid;
		}
		else {
			high = mid - 1;
		}
	}
	return answer + 1;
}

int main() {
	ios_base::sync_with_stdio(false);
	cin.tie(NULL);
	cin >> n;
	for (int i = 1; i <= n; i++) {
		cin >> arr[i];
	}
	sgInit(1, n, 1);
	int m;
	cin >> m;
	for (int i = 0; i < m; i++) {
		int a;
		cin >> a;
		if (a == 1) {
			int b, c;
			cin >> b >> c;
			update(1, n, 1, b, c);
			arr[b] += c;
		}
		else {
			int b;
			cin >> b;
			cout << binarySearch(b) << '\n';
		}
	}
	return 0;
}

---

시간을 최적화시키는 두 알고리즘의 조합이라 신선했다.