[백준 1168] 요세푸스 문제 2 (Platinum III) (Feat. PBDS) - 세그먼트 트리 없이 풀기

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

1168번: 요세푸스 문제 2

컴퓨터공학과에서 <자료구조>를 수강한 학생들이라면 std::set이 레드-블랙 트리로 구현되어 있다는 것을 알 것이다. 레드-블랙 트리는 Binary Search Tree다. 즉 원소들이 나름의 순서를 갖추고 정렬되어 있다는 것이다. 정렬되어 있다는 것은, 마치 Binary Search 하듯이 k번째 원소를 O(logn) 안에 접근할 수 있음을 함의한다. 그러나 그것을 하기 위해서는 rb-tree의 각 노드들이 subtree의 개수를 저장하고 있어야 한다. std::set은 이 기능을 제공하지 않는데, 아마 메모리 문제 때문인 듯하다. 결국 노드 하나당 4B를 더 저장해야 한다는 뜻이니까.

그런데 이 기능을 제공하는 자료구조가 있다! 공식 라이브러리는 아니고, g++의 확장 라이브러리인 policy-based data structures, 줄여서 PBDS에서. 이 라이브러리는 각종 자료구조를 지원하는데, 헤더파일 대충 읽어보니까 trie 등의 각종 편의기능을 제공하는 듯하다. 어쨌든 PBDS에서 제공하는 ordered_set이라는 자료구조는 std::set의 기능을 모두 지원하면서 다음 두 가지 기능을 추가적으로 지원한다.

• order_of_key(k) : k 미만의 원소의 개수를 반환한다.
• find_by_order(k) : k번째 원소를 가리키는 iterator를 반환한다. (당연히 index는 0부터 시작)

이 모든 기능이 O(logn)이므로 요세푸스 문제를 O(nlogn) 안에 풀 수 있다.

ordered_set의 실제 이름은 약간 복잡한데, 그건 그냥 밑의 코드(using문)를 참고하면 된다. 특이사항으로 __gnu_pbds::tree 클래스는 무슨 템플릿 인자를 5개씩이나 받는다. 이 템플릿 인자들이 뭘 하는지는 여기를 참고하면 된다.

#include <bits/stdc++.h>
#include <bits/extc++.h>
using ordered_set = __gnu_pbds::tree<int, __gnu_pbds::null_type, std::less<int>, __gnu_pbds::rb_tree_tag, __gnu_pbds::tree_order_statistics_node_update>;

int main() {
    int n, k; std::cin >> n >> k; if (std::cin.fail()) throw std::exception();

    ordered_set os;
    for (int i = 1; i <= n; ++i) {
        os.insert(i);
    }

    std::cout << '<';
    for (int _k = 0; not os.empty();) {
        _k += k - 1;
        if (_k >= os.size()) {
            _k %= static_cast<int>(os.size());
        }
        auto it = os.find_by_order(_k);
        std::cout << *it << (os.size() == 1 ? ">" : ", ");
        os.erase(it);
    } std::cout << std::endl;

    return 0;
}