[백준알고리즘] 1168번: 요세푸스 문제 2 -C++

[백준알고리즘] 1168번: 요세푸스 문제 2 -C++

1168번: 요세푸스 문제 2 (acmicpc.net)

1168번: 요세푸스 문제 2

예전에 Python으로 코테 문제를 풀었을 때 \(O(NK)\) 만에 통과하는 코드로 풀었다.

 

아래가 해당 글의 링크다.

[백준알고리즘] 1158, 1168번: 요세푸스 문제, 요세푸스 문제 2 -Python (tistory.com)

[백준알고리즘] 1158, 1168번: 요세푸스 문제, 요세푸스 문제 2 -Python

 

댓글로 통과하지 못한다는 것을 보고 찾아봤는데 최근 3개월, 5개월 안에 시간제한이 타이트하게 변경되어 그렇다는 내용이었다.

 

그래서 정말.. 아무리 생각해도 \(O(NK)\)로는 통과 불가능한데 그것만 주구장창 생각하다가 많은 시간도 보내고 많은 시간 초과를 경험했다...ㅎ..

 

끔찍하다

 

그래서 결국 찾아보다가.. 세그먼트 트리를 이용해야 한다는 것을 보고 다시금 세그먼트 트리(Segment Tree)를 공부하고 문제를 풀었다.

세그먼트 트리는 Crocus님의 블로그의 글을 통해 다시 공부했다. 

 

괜히 플레티넘 난이도가 아니기 때문에 세그먼트 트리에 대해 완벽히 이해하고 문제를 봐야한다.

글을 쓰고 있는 나조차 제대로 푼 건지도 모를 정도로 어렵다ㅠㅠ

다 쓰고 보니 되게 정신없이 문제를 풀어서 쓴 것 같은데.. 

 

잘못된 점이나 부족한 점 지적해주시면 감사하겠습니다

 

---

 

문제를 세그먼트 트리로 풀기 위해서는 아래의 과정을 거쳐야 한다.

1. 세그먼트 트리 생성
   1. 모든 리프 노드 1로 세팅
2. 환형으로 돌아가며 요세푸스 순열을 만들기 위해 필요한 값을 찾음
3. 필요한 값의 리프 노드를 0으로 초기화

위의 과정에서 2번을 제외하고는 모두 세그먼트 트리를 통해 쉽게 처리가 가능하다.

내가 이해가 어려웠던 부분도 2번 부분이었는데, 이 문제를 풀기 위해서는 2번을 나눠서 생각해야 한다.

 

 

예제 입출력을 통해 예를 들자면 {1, 2, 3, 4, 5, 6, 7}이라는 번호로 이루어진 환이 있다고 생각하고, 거기서 3번째마다 수를 하나씩 뽑아야 한다.

 

이때 리프 노드가 1 또는 0인 세그먼트 트리를 이용하면 1부터 임의의 번호 t까지의 누적합으로 1부터 t까지 몇 개의 숫자가 아직 요세푸스 순열에 사용되지 않았는지 알 수 있다.

가장 먼저 3을 요세푸스 순열에 사용했다면, {1, 2, x, 4, 5, 6, 7}로 남았다고 생각해보자. 3까지 살아있는 숫자의 개수는 {1, 2}로 2개이다. 여기서부터 앞으로 3개의 값을 더 찾아야 한다. 따라서 뒤이어 나오는 {4, 5, 6} 3개의 숫자를 추가적으로 세야 하며, 최종적으로 {1, 2, 4, 5, 6}으로 1부터 5개의 숫자가 사용되어야 한다는 것을 알 수 있다.

위의 말을 다르게 말한다면, 현재 뽑힌 숫자를 제외하고 3개의 숫자가 포함되어야 하기 때문에 현재 요세푸스 순열에 사용한 숫자 이후 부분 누적합이 3인 영역을 찾아야 한다는 것이다.

이 부분 누적합이 3이 되는 마지막 인덱스가 다음 요세푸스 순열에 추가될 숫자가 된다. 따라서 여기서 {4, 5, 6} 중 6이 마지막 인덱스이기 때문에 6이 요세푸스 순열에 3 다음으로 들어가게 된다.

 

즉, 이 과정도 두 단계로 나눠서 생각해야 한다.

1. 1부터 현재 지점까지의 누적합 S에 부분 누적합이 K가 추가되어, S+K의 누적합이 필요하다.
2. 순열을 만들기 위해 필요한 것은 누적합이 아닌, 환형으로 순차적으로 탐색했을 때 현재 지점부터 K 부분 누적합을 만드는 지점의 인덱스가 필요하다.

1번 과정은 환형인 점을 고려해 QuerySum 함수로 구현했다.

2번 과정은 환형인 점을 고려해 QuerySum을 통해 나온 누적합 지점을 만족하는 인덱스를 찾는 과정을 QueryIndex 함수로 구현했다.

 

#include <iostream>
#include <vector>
#include <cmath>

void Solve(void);

int main(void)
{
	std::ios::sync_with_stdio(false);
	std::cin.tie(NULL);
	std::cout.tie(NULL);
	
	Solve();
}

void MakeTree(std::vector<int>& segmentTree, int arrSize)
{
	int t = ceil(log2(arrSize));
	int treeSize = 1 << (t + 1);

	segmentTree.resize(treeSize);
}

int UpdateTree(int val, std::vector<int>& segmentTree, int nodeIndex, int arrStart, int arrEnd, int arrTarget)
{
	if (arrTarget < arrStart || arrEnd < arrTarget)
		return segmentTree[nodeIndex];
	
	if (arrStart == arrEnd)
	{
		if (arrStart == arrTarget)
			segmentTree[nodeIndex] = val;
		return segmentTree[nodeIndex];
	}

	int arrMid = (arrStart + arrEnd) / 2;
	segmentTree[nodeIndex] = UpdateTree(val, segmentTree, nodeIndex * 2, arrStart, arrMid, arrTarget)
						 + UpdateTree(val, segmentTree, nodeIndex * 2 + 1, arrMid + 1, arrEnd, arrTarget);
	return segmentTree[nodeIndex];
}

int QuerySum(std::vector<int>& segmentTree, int lowIndex, int highIndex, int nodeIndex, int arrStart, int arrEnd)
{
	if (arrEnd < lowIndex || highIndex < arrStart)
		return 0;

	if (lowIndex <= arrStart && arrEnd <= highIndex)
		return segmentTree[nodeIndex];

	int arrMid = (arrStart + arrEnd) / 2;
	return QuerySum(segmentTree, lowIndex, highIndex, nodeIndex * 2, arrStart, arrMid) 
		 + QuerySum(segmentTree, lowIndex, highIndex, nodeIndex * 2 + 1, arrMid + 1, arrEnd);
}

int QueryIndex(std::vector<int>& segmentTree, int val, int nodeIndex, int arrStart, int arrEnd)
{
	if (arrStart == arrEnd)
		return arrStart;

	int arrMid = (arrStart + arrEnd) / 2;
	if (val <= segmentTree[nodeIndex * 2])
		return QueryIndex(segmentTree, val, nodeIndex * 2, arrStart, arrMid);
	else
		return QueryIndex(segmentTree, val - segmentTree[nodeIndex * 2], nodeIndex * 2 + 1, arrMid + 1, arrEnd);
}

void Solve(void)
{
	int n, k;
	std::cin >> n >> k;

	std::vector<int> segmentTree;
	MakeTree(segmentTree, n);
	for (int i = 1; i <= n; i++)
		UpdateTree(1, segmentTree, 1, 1, n, i);

	std::vector<int> result;

	while (segmentTree[1])
	{
		int needSum = k;
		if (result.size())
			needSum = (QuerySum(segmentTree, 1, result.back(), 1, 1, n) + k ) % segmentTree[1];

		if (0 == needSum)
			needSum = segmentTree[1];

		result.push_back(QueryIndex(segmentTree, needSum, 1, 1, n));
		UpdateTree(0, segmentTree, 1, 1, n, result.back());
	}

	std::cout << "<";
	for (int i = 0; i < n; i++)
		std::cout << result[i] << (i == n - 1 ? ">" : ", ");
}

 

---

[참고]

Crocus

세그먼트 트리(Segment Tree)