Splay Tree - 3

작성일 | In

이전 글

Splay Tree 1: Rotate와 Splay
Splay Tree 2: insert, find, delete

kth element

약 1년 전에 올린 Splay Tree 2에서 예고한대로 Splay Tree에서 k번째 원소를 찾는 것을 구현해봅시다.

insert, find, delete는 C++의 set/map에서 지원해주기 때문에 저 기능들을 위해 굳이 BBST를 구현할 필요는 없었지만, k번째 원소를 찾는 기능은 없기 때문에 k번째 원소를 구하고 싶다면 BBST를 구현해야 합니다. pb_ds가 있긴 합니다...

각 노드를 루트로 하는 서브 트리의 크기를 알고 있어야 k번째 원소를 구할 수 있습니다. Node 구조체에 노드의 개수를 저장하는 변수를 추가합시다.

1
2
3
4
5
6
struct Node{
  Node* l;
  Node* r;
  Node* p;
  int key, cnt;
}

여기서 cnt 변수는 rotate가 될때마다 계속 바뀌기 때문에 cnt변수의 갱신을 담당할 update함수를 만듭시다.

1
2
3
4
5
void update(Node *x){
    x->cnt = 1;
    if(x->l) x->cnt += x->l->cnt;
    if(x->r) x->cnt += x->r->cnt;
}

rotate를 할때마다 cnt의 값이 변화하므로 rotate함수의 마지막 부분에서 update를 해주면 됩니다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
void rotate(node* x){
  node* p = x->p;
  node* b = NULL;
  if(!p) return;
  if(x == p->l){
    p->l = b = x->r;
    x->r = p;
  }else{
    p->r = b = x->l;
    x->l = p;
  }
  x->p = p->p; p->p = x;
  if(b) b->p = p;
  (x->p ? p==x->p->l ? x->p->l : x->p->r : tree) = x;
  update(p); update(x);
}

rotate가 끝난 시점에서, p가 x의 자식 노드가 되기 때문에 p를 먼저 갱신해줘야 합니다.

각 노드에 대해 cnt를 모두 구해놓았다면, 아래와 같이 k번째 원소를 구해줄 수 있습니다.

1
2
3
4
5
6
7
8
9
10
void kth(int k){
    Node* x = root;
    while(1){
        while(x->l && x->l->cnt > k) x = x->l;
        if(x->l) x -= x->l->cnt;
        if(!k--) break;
        x = x->r;
    }
    splay(x);
}

kth함수가 실행이 되면 k번째 노드가 루트로 올라오게 됩니다.

다음 글

다음 글에서는 위에서 구현한 update함수와 계속해서 사용하고 있는 splay함수를 이용해 구간 쿼리를 처리하는 방법에 대해 다룰 예정입니다.