Set 와 PBDS

우리가 흔히 쓰는 std::set 은,

특정값보다 작은 원소를 세거나, K번째 원소를 구하는 작업을 O(log N) 만에 할 수 없다.

하지만, PBDS (Policy Based Data Structure) 를 활용하면 가능하다.

 

물론, Set 을 안 쓰고도 할 수 있다.

특정값보다 작은 원소 개수는 세그먼트 트리로 알 수 있으며,

K번째 원소를 구하는 것은 세그먼트 트리 + 이분탐색으로 해결할 수 있다.

입력된 숫자 범위가 크더라도 좌표압축을 하거나, 동적 세그먼트 트리를 쓰면 될 것이다.

 

하지만, 위 기능들이 제공되는 Set 을 사용하면 훨씬 간단할 것이다.

 

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
#define ordered_set tree<int, null_type, less<int>, rb_tree_tag,tree_order_statistics_node_update>

 

우선 위 코드를 소스의 윗부분에 복붙하면 ordered_set 을 사용할 수 있다.

ordered_set 에서 insert() 와 erase() 는 std::set 와 동일하게 동작한다.

 

ordered_set OS; /// empty
    
OS.insert(1); /// 1
OS.insert(1); /// 1
OS.insert(4); /// 1 4
OS.insert(3); /// 1 3 4
OS.insert(2); /// 1 2 3 4
OS.insert(4); /// 1 2 3 4
OS.erase(2);  /// 1 3 4
OS.erase(2);  /// 1 3 4
OS.insert(1); /// 1 3 4
OS.erase(1);  /// 3 4

 

이제 ordered_set 의 중요한 기능 2가지를 살펴보자.

 

(Set 이름).order_of_key(NUM) : ordered_set 에서 NUM 보다 작은(미만의) 원소의 개수를 반환한다. 

(Set 이름).find_by_order(K) : ordered_set 에서 (K+1)번째 원소가 있는 iterator 을 반환한다. (K가 0이면 1번째)

 

실제 사용 예시를 보자.

 

1. order_of_key() 함수

/// OS : 13 32 34 45 48 59 61 68 71 78 87 128

cout << OS.order_of_key(10) << endl; /// 0

cout << OS.order_of_key(49) << endl; /// 5
cout << OS.order_of_key(48) << endl; /// 4

 

10 보다 작은 원소는 없으므로 0 이 반환된다.

49 보다 작은 원소는 13, 32, 34, 45, 48 으로, 5 가 반환된다.

48 보다 작은 원소는 13, 32, 34, 45 으로, 4가 반환된다.

 

2. find_by_order() 함수

/// OS : 13 32 34 45 48 59 61 68 71 78 87 128

cout << *(OS.find_by_order(0)) << endl; /// 13
cout << *(OS.find_by_order(7)) << endl; /// 68
cout << *(OS.find_by_order(100)) << endl; /// 0

 

find_by_order() 함수는 ordered_set 이 0번째 칸부터 채워진 1차원 배열이라고 생각하면 이해하기 편하다.

order == 0 이라는 것은 곧 1번째 수를 찾는 것이므로, 13 을 가리키는 iterator 이 반환된다.

order == 7 이라는 것은 곧 8번째 수를 찾는 것이므로, 68 을 가리키는 iterator 이 반환된다.

order == 100 이면 ordered_set 을 벗어나므로 유효하지 않은 iterator 이 반환된다. (*it == 0)

 

std::set 에서는 O(N) 으로 밖에 할 수 없었던 이 두 기능은 ordered_set 으로 O(log N) 만에 처리 가능하다.

---

+ 추가 내용 : multiset 과 pbds

이제 드는 의문점은, 과연 multiset 도 pbds 에서 구현이 될까? 이다.

실제로 Set 이 필요한 문제들은 multiset 인 경우도 많기 때문이다.

 

ordered_set 을 multiset 처럼 쓰려면 less<int> 를 less_equal<int> 로 바꾸기만 하면 된다.

 

#include <ext/pb_ds/assoc_container.hpp>
#include <ext/pb_ds/tree_policy.hpp>
using namespace __gnu_pbds;
#define ordered_set tree<int, null_type, less_equal<int>, rb_tree_tag,tree_order_statistics_node_update>

 

위 처럼 less_equal<int> 로 바꿔주자.

그러면 ordered_set 에 같은 원소를 여러번 삽입할 수 있다.

하지만 문제가 있다. erase() 함수가 제대로 동작하지 않는다. 아래 코드를 보자.

 

ordered_set OS; /// empty

OS.insert(1); OS.insert(1); OS.insert(1); /// 1 1 1
OS.erase(1); /// 1 1 1
OS.erase(1); /// 1 1 1

 

주석이 OS 의 상태를 나타내는 것인데, 1을 3번 insert 하는 것까지는 문제가 없다.

하지만 1을 erase 할 때 제대로 지워지지 않고 그대로 3개가 남아있게 된다.

 

따라서 우리는 어쩔 수 없이 직접 구현해주어야 한다.

 

void m_erase(ordered_set &OS, int val){
    int index = OS.order_of_key(val);
    ordered_set::iterator it = OS.find_by_order(index);
    OS.erase(it);
}

 

index 는 val 보다 작은 원소의 개수이므로, val 은 ordered_set 의 index + 1 번째 수일 것이다.

따라서 find_by_order(index) 를 하면 val 이 있는 정확한 위치를 알 수 있다.

 

이제 m_erase 함수가 잘 작동하는지 보자.

 

ordered_set OS; /// empty

OS.insert(1); OS.insert(1); OS.insert(1); /// 1 1 1
m_erase(OS, 1); /// 1 1
m_erase(OS, 1); /// 1

 

실제로 1이 1개씩 잘 줄어드는 것을 확인할 수 있다.

하지만, 이렇게 하면 multiset 에 없는 원소를 호출해도 잘 지워진다. 밑의 코드로 알 수 있다.

 

ordered_set OS; /// empty

OS.insert(1); OS.insert(1); OS.insert(1); /// 1 1 1
m_erase(OS, 5); /// 1 1 1
m_erase(OS, -6); /// 1 1

 

-6을 지웠는데도 1이 없어진다. 왜인지는 쉽게 알 수 있다.

이를 해결하기 위해서는, OS 에 해당 원소가 있는지부터 검사한 후 지워야 할 것이다.

이렇게 해서 만들어진 최종 m_erase() 함수는 아래와 같다.

 

void m_erase(ordered_set &OS, int val){
    int index = OS.order_of_key(val);
    ordered_set::iterator it = OS.find_by_order(index);
    if(*it == val) OS.erase(it);
}

 

erase() 를 실행하기 전, OS 에 val 이 없으면 삭제 작업은 하지 않는 명령을 추가하면 된다.

그러면 아래처럼 잘 작동한다. 

 

ordered_set OS; /// empty

OS.insert(1); OS.insert(1); OS.insert(1); /// 1 1 1
m_erase(OS, 5); /// 1 1 1
m_erase(OS, -6); /// 1 1 1
m_erase(OS, 1); /// 1 1

 

multiset ordered_set 은 find() 함수도 erase() 함수와 마찬가지로 제대로 작동이 안 되는 것 같다...

OS = {1, 1, 1} 에서 find(1) 을 하면 OS.end() 가 반환되기 때문이다.

 

따라서, multiset 에서 find() 함수는 m_erase() 처럼 find_by_order() 을 이용해야 할 것이다.

order_of_key() 로 그 원소보다 작은 원소 수를 찾고, find_by_order() 로 해당 위치를 알아낸 후,

우리가 찾는 원소랑 일치하는지 비교하면 된다.