BOJ 3653 : 영화 수집

알고리즘 문제 전체 링크

 

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

 

참고

- 구간 합 구하기 with 탑 다운 세그먼트 트리 (Top-Down Segment Tree)

- 구간 합 구하기 with 바텀 업 세그먼트 트리 (Bottom-Up Segment Tree)

 

 

6개의 영화가 있고 최대 M = 10번 영화를 본다고 가정하자.

그러면 세그먼트 트리를 아래와 같이 만들 수 있다.

마지막 leaf에서 10개는 비워두고, 6개의 영화를 순서대로 노드에 넣고 값을 1씩 증가시킨다.

영화를 10번 보기 때문에 앞의 노드를 그만큼 비워둔다.

 

영화가 순서대로 1, 2, 3, 4, 5, 6이 쌓여있는 상태다.

예를 들어 3번 영화를 본다고 가정하면, 3번 앞에는 1번과 2번 영화가 있기 때문에 2를 출력해야 한다.

이때 3번 앞에 있는 영화는 M + 1 번째 노드 ~ M + 3 - 1 번째 노드의 합이 된다.

 

구간의 합은 세그먼트 트리를 이용하면 간단히 구할 수 있으므로 문제를 풀 수 있다.

 

이제 3번 영화를 보았으므로, 앞으로 이동시킨다.

M + 3 = 13번째 노드는 -1을 하고 M = 10번째 노드를 1씩 증가시키면 트리는 아래와 같아진다.

3번 영화의 index는 13 → 10으로 바꾸고, M을 1 감소 시킨다. (M → 9)

 

만약에 4번 영화 앞에 있는 영화의 개수를 구하려면

M + 1번째 노드 ~ 14 - 1번째 노드의 합을 구하면 된다.

이때 M이 9로 변경되었고 10 ~ 13번째 노드의 합 = 3 을 구하게 된다. 

---

구현

 

tc가 여러 개이므로 tree를 0으로 초기화하는 코드가 필요하다.

for (int i = 1; i <= treeSize; i++) tree[i] = 0;

 

마지막 leaf에 저장되는 index는 M을 더해야하므로 아래와 같이 인덱스를 저장하고 트리에 업데이트 한다.

for (int i = 1; i <= N; i++) movieIndex[i] = M + i;
for (int i = 1; i <= N; i++) update(M + i, 1);

 

위에 설명한대로 현재 영화를 선택한 후, 구간의 합을 구한다.

그리고 index와 M(= queryCount)을 갱신한다.

for (int i = 1; i <= M; i++)
{
	int movie;

	scanf("%d", &movie);

	printf("%d ", getSum(queryCount + 1, movieIndex[movie] - 1));

	update(movieIndex[movie], -1);
	movieIndex[movie] = queryCount;
	queryCount--;
	update(movieIndex[movie], 1);
}

 

전체 코드는 다음과 같다.

#include <stdio.h>

int T, N, M;
int tree[524288];
int movieIndex[100100];
int start;

void update(int index, int diff)
{
	index += start;

	while (index > 1)
	{
		tree[index] += diff;
		index /= 2;
	}
}

int getSum(int left, int right)
{
	int ans = 0;

	left += start;
	right += start;

	while (left < right)
	{
		if (left % 2 == 0) left /= 2;
		else
		{
			ans += tree[left];
			left += 1;
			left /= 2;
		}

		if (right % 2 == 1) right /= 2;
		else
		{
			ans += tree[right];
			right -= 1;
			right /= 2;
		}
	}

	if (left == right) ans += tree[left];

	return ans;
}

int main()
{
	scanf("%d", &T);
	for (int tc = 0; tc < T; tc++)
	{
		int n;

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

		for (n = 1; (1 << n) <= (N + M); n++);
		start = 1 << n;
		start--;

		int treeSize = (start + 1) * 2;
		for (int i = 1; i <= treeSize; i++) tree[i] = 0;
		for (int i = 1; i <= N; i++) movieIndex[i] = M + i;
		for (int i = 1; i <= N; i++) update(M + i, 1);

		int queryCount = M;
		for (int i = 1; i <= M; i++)
		{
			int movie;

			scanf("%d", &movie);

			printf("%d ", getSum(queryCount + 1, movieIndex[movie] - 1));

			update(movieIndex[movie], -1);
			movieIndex[movie] = queryCount;
			queryCount--;
			update(movieIndex[movie], 1);
		}

		putchar('\n');
	}

	return 0;
}

---

탑 다운 방식은 다음과 같다. (메모리를 조금 더 넉넉하게 잡았다.)

#include <stdio.h>

int T, N, M;
int tree[524288 * 2];
int arr[100100];
int movieIndex[100100];

int getSum(int left, int right, int a, int b, int node)
{
	if (b < left || right < a) return 0;
	if (a <= left && right <= b) return tree[node];

	int mid = (left + right) / 2;

	return getSum(left, mid, a, b, node * 2) + getSum(mid + 1, right, a, b, node * 2 + 1);
}

void update(int left, int right, int node, int index, int diff)
{
	if (index < left || right < index) return;

	tree[node] += diff;

	if (left != right)
	{
		int mid = (left + right) / 2;

		update(left, mid, node * 2, index, diff);
		update(mid + 1, right, node * 2 + 1, index, diff);
	}
}

int main()
{
	scanf("%d", &T);
	for (int tc = 0; tc < T; tc++)
	{
		scanf("%d %d", &N, &M);

		int treeSize = (N + M) * 3;
		for (int i = 1; i <= treeSize; i++) tree[i] = 0;
		for (int i = 1; i <= N; i++) movieIndex[i] = M + i;
		for (int i = 1; i <= N; i++) update(1, N + M, 1, M + i, 1);

		int queryCount = M;
		for (int i = 1; i <= M; i++)
		{
			int movie;

			scanf("%d", &movie);

			printf("%d ", getSum(1, N + M, queryCount + 1, movieIndex[movie] - 1, 1));

			update(1, N + M, 1, movieIndex[movie], -1);
			movieIndex[movie] = queryCount;
			queryCount--;
			update(1, N + M, 1, movieIndex[movie], 1);
		}

		putchar('\n');
	}

	return 0;
}