백준 1168 요세푸스 문제 2

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

 

요세푸스 문제는 매번 남은 수중 (현재 위치 + K) % 수의 개수 번째 수를 찾는 문제이다. 이는 세그먼트트리로 K번째 수를 찾는 연산과 존재하던 수를 지우는 연산을 구현해 풀 수 있다.

 

#pragma GCC optimize("Ofast")
#include <bits/stdc++.h>
using namespace std;

int res = 1;

vector<int> tree;

void init(int t1, int t2, int idx){
    if (t1 == t2){
        tree[idx] = 1;
        return;
    }
    int mid = (t1 + t2) >> 1;
    init(t1, mid, idx << 1);
    init(mid + 1, t2, idx << 1 | 1);
    tree[idx] = tree[idx << 1] + tree[idx << 1 | 1];
}

void init(int N){
    tree.resize(res);
    init(0, N - 1, 1);
}

void update(int t1, int t2, int q, int idx){
    if (q < t1 || t2 < q) return;
    if (t1 == t2){
        tree[idx]--;
        return;
    }
    int mid = (t1 + t2) >> 1;
    update(t1, mid, q, idx << 1);
    update(mid + 1, t2, q, idx << 1 | 1);
    tree[idx] = tree[idx << 1] + tree[idx << 1 | 1];
}

int query(int t1, int t2, int K, int idx){
    if (t1 == t2) return t1;
    int mid = (t1 + t2) >> 1;
    if (K < tree[idx << 1]) return query(t1, mid, K, idx << 1);
    else return query(mid + 1, t2, K - tree[idx << 1], idx << 1 | 1);
}

int main(){
    ios_base::sync_with_stdio(false);
    cin.tie(NULL);
    
    int N, K;
    cin >> N >> K;
    
    while (res < N << 1){
        res <<= 1;
    }
    init(N);
    vector<int> ans;
    int no = N - 1;
    
    for (int i = N; i >= 1; i--){
        no = (no + K) % i;
        int qu = query(0, N - 1, no, 1);
        no--;
        update(0, N - 1, qu, 1);
        ans.push_back(qu + 1);
    }
    
    cout << '<';
    for (int i = 0; i < N; i++){
        cout << ans[i];
        if (i + 1 == N) cout << ">\n";
        else cout << ", ";
    }
}