BOJ 1168: 요세푸스 문제 2

www.acmicpc.net/problem/1168

1168번: 요세푸스 문제 2

 

$N$, $K$가 주어졌을 때 요세푸스 순열을 순서대로 출력하는 문제이다.

순열을 간단히 요약하자면 $K$칸씩 건너뛰면서 해당 위치에 있는 번호를 제거하는 순열이다. 단 중간에 빠진 사람을 카운트해서는 안된다.

 

즉 어떤 구간의 합이 $K$인지를 파악하고 딱 그만큼이 되는 부분으로 이동할 수 있게끔 만드는 것이 핵심이라 생각했기 때문에 segment tree를 이용하여 $K$ 번째 수를 구하는 문제와 유사하다고 생각하여 접근하였다.

 

Seg tree를 이용하여 $K$ 번째 수를 구하는 것은 이분 탐색과 비슷한 원리로 구현할 수 있다. 먼저 현재 노드에서 왼쪽 자식($V_l$이라고 하자.)과 $K$를 비교한다. $V_l \leq K$를 만족하는 경우 찾고자 하는 수가 왼쪽에 존재하는 것이고 그렇지 않은 경우 오른쪽에 존재한다. 여기서 오른쪽 자식 노드로 탐색을 진행하는 경우에는 오른쪽에서 $K - V_l$ 번째 수를 찾는 것이므로 이를 고려하여 구현해야 한다.

 

이를 이용한다면 생각보다 쉽게 구현될거라 생각했지만 몇 가지 문제가 존재했다.

1. 현재 위치가 1이 아닌 경우
2. 현재 위치를 기준으로 뒤에 남아있는 숫자가 $K$보다 작은 경우
3. 2번과 비슷하지만 현재 위치가 1이고 현재 남아있는 숫자가 $K$보다 작은 경우

첫 번째 문제점은 단순히 현재 위치를 가리키는 변수($cur$)를 추가한 후, 해당 위치를 기준으로 구간 $[1, cur)$의 합을 구하여 왼쪽 자식에서 항상 이 값을 빼주는 방법을 이용하여 해결할 수 있다. 이때 왼쪽 자식으로 이동하는 경우에는 이 값을 유지해줘야만 하며 오른쪽 자식으로 넘어가는 경우 이 값을 0으로 설정해주면 된다.

 

두 번째, 세 번째 문제점은 유사하다. 두 번째 문제의 경우$(cur, N)$의 합을 $K$에서 빼준 후 첫 번째 위치로 이동시켜주면 쉽게 해결할 수 있다.

 

세 번째 문제점은 맨 처음 단순히 $K$를 구간 $[1, N]$의 합으로 나눈 나머지를 넣는 방식으로 해결하려고 했으나 나머지가 $0$인 경우에 문제가 발생하여 단순히 이를 예외처리 해줌으로써 해결하였다.

 

 

#include<bits/stdc++.h>
using namespace std;
typedef pair<int, int> P;
typedef long long ll;
struct Segment {
	vector<ll> tree;
	vector<ll> seq;
	Segment(int n) {
		++n;
		seq.resize(n);
		int log = ceil(log2(n));
		int t = (1 << (log + 1));
		tree.resize(t);
	}

	void init(int node, int s, int e) {
		if (s == e) {
			tree[node] = seq[s];
		}
		else {
			init(node * 2, s, (s + e) / 2);
			init(node * 2 + 1, (s + e) / 2 + 1, e);
			tree[node] = tree[node * 2] + tree[node * 2 + 1];
		}
	}

	ll query1(int node, int s, int e, int i, int j) {
		if (e < i || s > j) return 0;
		if (i <= s && e <= j) return tree[node];
		return query1(node * 2, s, (s + e) / 2, i, j) + query1(node * 2 + 1, (s + e) / 2 + 1, e, i, j);
	}
	int query2(int node, int s, int e, int num, int left) {
		if (s == e)
			return s;
		int l = tree[node * 2] - left;
		if (l < num)
			return query2(node * 2 + 1, (s + e) / 2 + 1, e, num - l, 0);
		else
			return query2(node * 2, s, (s + e) / 2, num, left);
	}

	void update(int node, int s, int e, int i, ll u) {
		if (i < s || i > e) return;
		tree[node] += u;
		if (s != e) {
			update(node * 2, s, (s + e) / 2, i, u);
			update(node * 2 + 1, (s + e) / 2 + 1, e, i, u);
		}
	}
};
int main() {
	int n, k; scanf("%d %d", &n, &k);
	Segment s(n);
	for (int i = 1; i <= n; ++i) {
		s.seq[i] = 1;
	}
	s.init(1, 1, n);
	printf("<");
	int cur = 1;
	int move = k;
	for(int i = 0; i<n; ++i){
		int l = s.query1(1, 1, n, 1, cur - 1);
		int r = s.query1(1, 1, n, cur, n);
		if (cur == 1){
			if (s.tree[1] < move) {
				move %= s.tree[1];
				if (move == 0) move = s.tree[1];
				--i;
				continue;
			}
		}
		if (r < move) {
			move -= r;
			--i;
			cur = 1;
			continue;
		}
		int ans = s.query2(1, 1, n, max(move, 1), l);
		s.update(1, 1, n, ans, -1);
		printf("%d%s", ans, i == n - 1 ? "" : ", ");
		cur = ans;
		move = k;
	}
	printf(">");
}