Featured image of post C++ STL PBDS (Policy based data structures)

C++ STL PBDS (Policy based data structures)

C++ STL PBDS ordered_set / ordered_multi_set 사용법과 boj 1572 중앙값 풀이

ordered_set이란?

g++에서 추가된 자료구조로, std::set과 유사하지만, 아래 두 연산을 $\mathcal{O}(\log N)$만에 수행할 수 있다.

  • order_of_key(k) : k보다 작은 원소의 개수

  • find_by_order(k) : 오름차순으로 정렬했을 때 k번째 원소 (zero-based)

집합 내부에서 원소의 추가/삭제가 빈번하게 이루어지며, 수시로 k번째 수를 찾아야하는 경우에 유용한 자료구조이다.

1
2
3
4
5
6
7
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
typedef tree<int, null_type, less<>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;

ordered_set os;

다만, 삭제하는 연산은 아래와 같이 따로 구현해주어야 한다.

1
2
3
4
5
void erase(ordered_set &os, int value) {
    size_t index = os.order_of_key(value);
    auto iter = os.find_by_order(index);
    os.erase(iter);
}

ordered_set 사용 예제

order_of_key의 경우 value를 리턴하고, find_by_order의 경우 iterator를 리턴한다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
#include <iostream>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
typedef tree<int, null_type, less<>, rb_tree_tag, tree_order_statistics_node_update> ordered_set;

void erase(ordered_set &os, int value) {
    size_t index = os.order_of_key(value);
    auto iter = os.find_by_order(index);
    os.erase(iter);
}

int main() {
    ordered_set os;

    os.insert(5);
    os.insert(4);
    os.insert(3);
    os.insert(3);
    os.insert(2);
    os.insert(1);

    for (int i : os) cout << i << ' ';  // 1 2 3 4 5
    cout << os.order_of_key(2) << endl;  // 1
    cout << *os.find_by_order(2) << endl;  // 3

    cout << os.size() << endl;  // 5
    erase(os, 2);
    cout << os.size() << endl;  // 4

    for (int i : os) cout << i << ' ';  // 1 3 4 5
    cout << os.order_of_key(2) << endl;  // 1
    cout << *os.find_by_order(2) << endl;  // 4
}

ordered_multi_set 사용 예제

중복 값을 허용한다.

less<> 대신 less_equal<>를 넣어주면 된다.

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
#include <iostream>
#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace std;
using namespace __gnu_pbds;
typedef tree<int, null_type, less_equal<>, rb_tree_tag, tree_order_statistics_node_update> ordered_multi_set;

void erase(ordered_multi_set &oms, int value) {
    size_t index = oms.order_of_key(value);
    auto iter = oms.find_by_order(index);
    oms.erase(iter);
}

int main() {
    ordered_multi_set oms;

    oms.insert(5);
    oms.insert(4);
    oms.insert(3);
    oms.insert(3);
    oms.insert(2);
    oms.insert(1);

    for (int i : oms) cout << i << ' ';  // 1 2 3 3 4 5
    cout << oms.order_of_key(2) << endl;  // 1
    cout << *oms.find_by_order(2) << endl;  // 3

    cout << oms.size() << endl;  // 6
    erase(oms, 2);
    cout << oms.size() << endl;  // 5


    for (int i : oms) cout << i << ' ';  // 1 3 3 4 5
    cout << oms.order_of_key(2) << endl;  // 1
    cout << *oms.find_by_order(2) << endl;  // 3
}

BOJ 1572 중앙값

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

세그트리/이분탐색 문제이지만, PBDS를 사용하면 쉽게 풀 수 있다.

문제 풀이

온도 값이 1초마다 추가되며, 최근 K초 까지 온도의 중앙값을 다 더한 값을 구하는 문제이다.

매 초마다 ordered_multi_set에 값을 추가해 준다.

k초 이전에는 값 입력만 받고, k초부터는 중앙값을 구해준다. find_by_order()는 zero-based임에 유의하자.

우리는 최근 k초까지의 값에만 관심 있기 때문에, k초 이후에는 k초 이전의 값은 삭제해 준다.

중복값이 발생하므로 less_equal<>를 사용해 주자. 정답은 $2^{63}-1$ 이하이므로, int범위를 초과한다.