BOJ 1321 군인

문제

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

전쟁에는 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줄에 걸쳐 명령이 주어진다.

출력

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

예제 입력 1

5

1 4 3 5 2

5

1 2 -3

1 4 2

2 5

1 2 4

2 5

예제 출력 1

3

2

풀이

이 문제는 파라매트릭 바이너리 서치와 세그먼트 트리를 이용하는 문제이다. 전체적인 알고리즘이 BOJ 1849 문제와 상당히 유사하다.

flappybird.tistory.com/9

BOJ 1849 순열

우선 초기 상태의 부대 인원으로 세그먼트 트리를 만든다. 그리고 세그먼트 트리의 각 노드에는 해당하는 영역의 합을 저장한다. 이를 통해 임의의 구간의 구간합을 lgN 만에 구할 수 있다. 또, 하나의 리프를 갱신하는 과정(1 i a 형태의 쿼리)도 lg N만에 실행할 수 있다.

 

이때 두 번째 쿼리를 처리하기 위해 파라매트릭 바이너리 서치를 사용하면 된다. 1부터 x까지 부대의 합이 i를 넘어가는(초과) 최소 x값을 바이너리 서치로 구현할 수 있다.

소스코드

#include <stdio.h>
#include <math.h>

#define MAX 2000000

int tree[MAX];
int l[MAX];
int r[MAX];

int s;

void init(int x)
{
	if (s <= x && x < s * 2)
	{
		l[x] = x - s + 1;
		r[x] = x - s + 1;
		return;
	}
	init(x * 2);
	init(x * 2 + 1);
	l[x] = l[x * 2];
	r[x] = r[x * 2 + 1];
	tree[x] = tree[x * 2] + tree[x * 2 + 1];
}

int sumquery(int loc, int x, int y)
{
	if (x == y)
	{
		return tree[x + s - 1];
	}
	if (l[loc] == x && r[loc] == y)
	{
		return tree[loc];
	}
	if (r[loc * 2] >= y)
	{
		return sumquery(loc * 2, x, y);
	}
	if (l[loc * 2 + 1] <= x)
	{
		return sumquery(loc * 2 + 1, x, y);
	}
	return sumquery(loc * 2, x, r[loc * 2]) + sumquery(loc * 2 + 1, l[loc * 2 + 1], y);
}

void update(int x, int a)
{
	if (!x)
	{
		return;
	}
	tree[x] += a;
	update(x / 2, a);
}

int findquery(int a, int N)
{
	int l, r, mid;
	
	l = 1;
	r = N;
	while (l < r)
	{
		mid = (l + r) / 2;
		if (sumquery(1, 1, mid) >= a)
		{
			r = mid;
		}
		else
		{
			l = mid + 1;
		}
	}
	return r;
}

int main(void)
{
	int N;
	scanf("%d", &N);

	s = 1 << (int)ceil(log2(N));

	int i;
	for (i = 1; i <= N; i++)
	{
		scanf("%d", &tree[i + s - 1]);
	}

	init(1);

	int M;
	scanf("%d", &M);

	int q, a, b;
	for (i = 1; i <= M; i++)
	{
		scanf("%d", &q);
		if (q == 1)
		{
			scanf("%d %d", &a, &b);
			update(a + s - 1, b);
		}
		else
		{
			scanf("%d", &a);
			printf("%d\n", findquery(a, N));
		}
	}
	
	return 0;
}