Splay Tree - 4

작성일 | In

이전 글

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

Splay Tree를 이용한 배열의 표현

rotate연산과 splay연산을 이용해 아무리 트리의 모양을 바꾼다고 해도 중위 순회를 하면 항상 똑같은 순서로 탐색하게 됩니다. BST이기 때문에 항상 키의 오름차순을 유지합니다.
이 성질을 이용해, 배열의 인덱스를 키로 생각해주면 배열의 각 원소를 splay tree의 노드에 저장할 수 있습니다.

이 글에서는 구간합 혹은 RMQ를 하는 방법과 특정 구간을 뒤집는 방법을 다룰 것입니다.

Splay 연산의 활용

Splay Tree는 이름 그대로 Splay 연산을 이용해 트리의 균형을 맞춥니다.
Splay연산을 이용해 원하는 노드를 트리의 루트로 올릴 수 있습니다. 이 연산을 이용해 세그먼트 트리처럼 선형 자료들을 관리해줄 수 있습니다.

이번 글에서는 splay연산을 이용해 어떤 구간 [s, e]를 한 노드에 모아주는 것에 집중해볼 것입니다.

구간 모으기

구간 [s, e]를 한 노드로 모아주는 방법을 먼저 알아봅시다.
s-1번째 노드를 splay해주면 아래 그림과 같이 됩니다.

e+1번째 노드를 s-1번째 노드의 오른쪽 자식으로 올려주면 아래 그림과 같이 됩니다.

이렇게 하면 루트 노드의 오른쪽 자식의 왼쪽 서브트리에 [s, e]구간이 모이게 됩니다.

splay함수를 약간 수정해서 어떤 노드를 다른 노드의 자식으로 붙이는 것을 만들어줍시다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
void splay(Node *x, Node *g = nullptr){
  Node* y;
  while(x->p != g){
      Node* p = x->p;
      if(p->p == g){
          rotate(x); break;
      }
      auto pp = p->p;
      if((p->l == x) == (pp->l == p)){
          rotate(p); rotate(x);
      }else{
          rotate(x); rotate(x);
      }
  }
  if(!g) root = x;
}

구간 [s, e]를 모으는 함수 gather는 아래와 같이 간단하게 구현할 수 있습니다.

1
2
3
4
5
6
7
Node* gather(int s, int e){ //gather [s, e]
    kth(e+1);
    auto tmp = root;
    kth(s-1);
    splay(tmp, root);
    return root->r->l;
}

백준2042 구간 합 구하기

k번째 원소를 구할 때 cnt를 갱신하기 위해 사용했던 update를 조금만 수정하면 왼쪽 서브트리와 오른쪽 서브트리의 합도 구해줄 수 있습니다.

1
2
3
4
5
6
7
8
9
10
11
12
void update(Node *x){
    x->sz = 1;
    x->sum = x->v;
    if(x->l){
        x->sz += x->l->sz;
        x->sum += x->l->sum;
    }
    if(x->r){
        x->sz += x->r->sz;
        x->sum += x->r->sum;
    }
}

지금까지 한 것을 모두 합치면 2042 구간 합 구하기를 풀 수 있습니다.
구현의 편의를 위해 왼쪽과 오른쪽에 dummy node를 하나씩 추가했습니다.
코드

백준10999 구간 합 구하기2

이제 lazy propagation을 추가해봅시다.
노드 구조체에 lazy를 저장할 변수를 추가합시다.

1
2
3
4
5
struct Node{
    ~~
    ll sum, lazy;
    ~~
};

lazy를 뿌려주는 함수 push도 구현해줍시다.

1
2
3
4
5
6
7
8
9
10
11
12
void push(Node *x){
    x->v += x->lazy;
    if(x->l){
        x->l->lazy += x->lazy;
        x->l->sum += x->l->sz * x->lazy;
    }
    if(x->r){
        x->r->lazy += x->lazy;
        x->r->sum += x->r->sz * x->lazy;
    }
    x->lazy = 0;
}

lazy는 자식으로 내려갈 때마다 뿌려줘야 합니다. kth함수를 약간 수정해줍시다.

1
2
3
4
5
6
7
8
9
10
11
12
13
14
void kth(int k){ //1-based
    auto now = root;
    push(now);
    while(1){
        while(now->l && now->l->sz > k){
            now = now->l; push(now);
        }
        if(now->l) k -= now->l->sz;
        if(!k) break; k--;
        now = now->r;
        push(now);
    }
    splay(now);
}

지금까지 한 것을 잘 조합하면 [10999 구간 합 구하기2]도 풀 수 있습니다.
코드

구간 뒤집기

구간을 뒤집는 것은 왼쪽 서브트리와 오른쪽 서브트리를 swap해주는 것을 재귀적으로 처리해주면 됩니다.
재귀적으로 처리하는 것은 느리기 때문에, flip여부를 lazy propagation 해주면 됩니다.

flip여부를 나타내기 위해 노드 구조체에 변수를 하나 만듭니다.

1
2
3
4
5
struct Node{
    ~~
    bool flip;
    ~~
};

flip여부를 전파하기 위한 push함수를 만들어줍시다.

1
2
3
4
5
6
7
void push(Node *x){
    if(!x->flip) return;
    swap(x->l, x->r);
    if(x->l) x->l->flip = !x->l->flip;
    if(x->r) x->r->flip = !x->r->flip;
    x->flip = 0;
}

여기까지 했다면, flip을 하는 것은 매우 간단합니다.

1
2
3
4
void flip(int s, int e){
    Node* x = gather(s, e);
    x->flip = !x->flip;
}