[백준/C++] #1168 요세푸스 문제 2 (Segment Tree)

문제

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

: 큐를 이용여 풀리는 가장 기본적인 요세푸스 문제

 

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

: 세그먼트 트리를 이용하여 시간복잡도를 줄여야 하는 요세푸스 문제

개념

: Segment Tree

 

코드 + 풀이

1. 큐를 이용한 풀이 (1158번)

: 특정 "번째"의 사람이 올 때까지 큐의 원소를 push, pop 반복한다. 

#include<iostream>
#include<queue>
using namespace std;

int n, k;
queue<int>q;

int main(void)
{
	int i;
	cin >> n >> k;
	for (i = 1; i <= n; i++)
		q.push(i);	//큐에 n까지 삽입
	cout << "<";
	while (q.size()!=0)	//큐가 빌 때까지 반복
	{
		for (i = 1; i < k; i++)
		{
			q.push(q.front());	//k-1번째 원소를 맨 뒤에 삽입
			q.pop(); // 맨 앞에 있던 원소는 삭제
		}
		cout << q.front(); //k번째 원소가 맨 앞에 오게 됨.
		if (q.size() != 1)
			cout << ", ";
		q.pop();//k번째 원소 출력 후 삭제
	}
	cout << ">";
	return 0;
}

 

2. 세그먼트 트리 이용한 풀이 (1168번)

: 요세푸스 문제의 핵심은 몇 번째에 있는 사람을 찾고, 해당 사람을 제거한 뒤 다시 특정 번째에 있는 사람을 찾을 수 있도록 업데이트 하는 것이다. 세그먼트 트리는 부분합을 구하는데에 사용되는데, 백준의 ' 2517번 달리기' 문제와 같이 트리를 통해 특정 조건의 명수를 구할 수 있다. 먼저 트리를 만들고(init), k번째 노드 번호를 찾고(query), 구해진 노드 번호를 통해 해당 노드를 제거한 후에 트리를 업데이트(update) 하면 된다. 

* idx: 찾아야하는 노드의 앞에서부터의 "번째"

* num: 찾아야 하는 노드의 트리 "인덱스"

#include <iostream>
#include <vector>
#include <queue>
#include <algorithm>
// 1168
/*
요세푸스 문제인데 시간 복잡도가 훨씬 작아야 한다. 
원래는 큐 사용 O(N^2) -> 시간 복잡도 줄이기 위해 세그먼트 트리 이용 O(N logM) 
*/
using namespace std;

int N, K;
vector<int> tree;

// 세그먼트 트리 만들기
int init(int node, int start, int end) {
    if (start == end) {
        return tree[node] = 1; // 명수니까 1
    }
    else {
        int mid = (start + end) / 2;
        return tree[node] = init(node * 2, start, mid) + init(node * 2 + 1, mid + 1, end);
    }
}

// 특정 번째에 있는 노드 번호 찾기
int query(int node, int start, int end, int order) {
    if (start == end) return start;
    int mid = (start + end) / 2;

    if (order <= tree[2 * node]) {
        return query(2 * node, start, mid, order);
    }
    else {
        return query(2 * node + 1, mid + 1, end, order - tree[2 * node]); // 중요! 남은 크기 계산 -> 다시 연산
    }
}

// 노드 번호를 통해 노드 제거 후 트리 업데이트
void update(int node, int start, int end, int del) {
    tree[node]--;
    if (start == end) return;
    int mid = (start + end) / 2;
    if (del <= mid) {
        update(2 * node, start, mid, del);
    }
    else {
        update(2 * node + 1, mid + 1, end, del); 
        // 해당 값 자체를 가지고 연산을 하는게 아니기에 반환값 void로 설정
    }
}

int main() {
    ios::sync_with_stdio(false); cin.tie(NULL); cout.tie(NULL);
    cin >> N >> K;

    tree.resize(4 * N); // 4로 해서 크기 충분하게 조정하자. 3으로 하면 out of bounds 런타임 에러 뜸

    init(1, 1, N); // 트리 생성

    int idx = 0; // 순서 계산 잘해야 함
    cout << "<";

    for (int i = 0; i < N; i++) {
        int size = N - i; // 남아있는 사람의 명수
        idx = (idx + K - 1) % size; // 삭제해야 할 사람의 앞에서부터의 번째

        int num = query(1, 1, N, idx + 1); // 삭제해야 할 노드 번호 찾기
        if (1 != size) { // 2명 이상 남았다면, 
            cout << num << ", ";
        }
        else { // 1명만 남았다면
            cout << num;
        }

        update(1, 1, N, num); // 노드 번호로 삭제, 후처리
    }

    cout << ">" << endl;
    return 0;
}