Persistent Segment Tree

작성일 | In

서론

어떤 자료구조가 “persistent하다”라는 것은, 상태의 변화가 생기더라도 과거 상태를 모두 보존하고 있다는 것을 의미합니다.
Persistent Segment Tree(PST)는 세그먼트 트리의 갱신 과정을 모두 저장하고 있는 자료구조입니다.

BOJ16978 수열과 쿼리 22

문제 링크

공간 복잡도를 신경쓰지 않고 바로 떠오르는 풀이를 생각해보면 다음과 같습니다.

  • 1번 쿼리 : 갱신된 정보를 반영한 새로운 세그먼트 트리를 만듬
  • 2번 쿼리 : k번째 세그먼트 트리로 가서 구간합 쿼리

공간 복잡도는 O(QN)이겠지요.
이 문제는 트리의 정보가 갱신되는 각 시점마다 세그먼트 트리를 만들어서 이전 정보를 모두 들고 있어야 합니다. 공간 복잡도를 줄일 방법을 생각해봅시다.

1번 쿼리로 인해 값이 바뀌는 정점은 O(log N)개 뿐입니다. 나머지는 값이 그대로 유지됩니다.
이러한 특성을 이용해서 1번 쿼리가 주어질 때마다 값이 바뀌는 O(log N)개의 정점은 새로 만들어주고, 나머지 정점들은 기존 세그먼트 트리의 것을 재활용할 수 있습니다.

세그먼트 트리에서 빨간색 정점들의 값이 바뀐다고 합시다.

이때 새로운 세그먼트 트리를 전부 만드는 것이 아닌, 값이 바뀌는 O(log N)개의 정점만 새로 만들어 주는 것이 PST의 핵심 아이디어입니다.

O(log N)개의 정점을 제외하면 값이 바뀌지 않으므로 기존 정점을 재활용할 수 있습니다. 위 그림에서는 점선으로 연결 정보를 표시했습니다.

이 상태에서 파란색 정점들을 업데이트를 한다고 생각해봅시다.

이번에도 O(log N)개의 정점만 새로 만들어주고, 나머지는 재활용을 해주면 됩니다.

위 그림에서 회색 루트 정점이 0번 세그먼트 트리의 루트가 되고 분홍색, 파란색이 각각 1번, 2번 세그먼트 트리의 루트가 됩니다.

초기에 O(N)개의 정점을 생성하고, 각 업데이트마다 O(log N)개의 정점을 생성하므로 총 공간 복잡도는 O(N + Q log N)입니다.

구현

Dynamic Segment Tree를 알고 있다면, PST는 쉽게 구현할 수 있습니다.

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
37
38
39
40
41
42
typedef long long ll;
struct Node{
    Node *l, *r;
    ll v;
    Node(){ l = r = NULL; v = 0; }
};

//root[i] = i번째 세그먼트 트리의 루트
Node *root[101010]; //root[0] 할당 필수
int arr[101010]; //초깃값

void build(Node *node, int s, int e){ //0번 트리 생성
    if(s == e){
        node->v = arr[s]; return;
    }
    int m = s + e >> 1;
    node->l = new Node(); node->r = new Node();
    build(node->l, s, m); build(node->r, m+1, e);
    node->v = node->l->v + node->r->v;
}
void add(Node *prv, Node *now, int s, int e, int x, int v){
    if(s == e){
        now->v = v; return;
    }
    int m = s + e >> 1;
    if(x <= m){ //왼쪽 자식에 업데이트 하는 경우
        //왼쪽 자식은 새로운 정점 생성, 오른쪽 자식은 재활용
        now->l = new Node(); now->r = prv->r;
        add(prv->l, now->l, s, m, x, v);
    }else{
        //오른쪽 자식은 새로운 정점 생송, 왼쪽 자식은 재활용
        now->l = prv->l; now->r = new Node();
        add(prv->r, now->r, m+1, e, x, v);
    }
    now->v = now->l->v + now->r->v;
}
ll query(Node *node, int s, int e, int l, int r){
    if(r < s || e < l) return 0;
    if(l <= s && e <= r) return node->v;
    int m = s + e >> 1;
    return query(node->l, s, m, l, r) + query(node->r, m+1, e, l, r);
}

BOJ7469 K번째 수

문제 링크

세그먼트 트리를 이용해 전체 구간에서 K번째로 작은 수를 찾는 것은 잘 알려져 있습니다. 이것을 이용해서 풀이를 찾아봅시다.

PST는 누적합의 개념도 갖고 있습니다. i번째 세그먼트 트리는 [0, i]번째 갱신을 모두 반영한 세그먼트 트리입니다. 그러므로 E번째 세그먼트 트리의 값에서 S-1번째 세그먼트 트리의 값을 빼주면, [S, E]번째 갱신으로 인해 값이 어떻게 바뀌었는지 알 수 있습니다.

우리가 구해야 하는 것은 [i, j] 구간에서 K번째로 작은 수입니다.
PST에 1부터 N번째 원소까지 차례대로 추가를 한 다음, j번째 세그먼트 트리에서 i-1번째 세그먼트 트리를 뺀 트리에서 K번째 원소를 찾으면 됩니다.

K번째 원소를 찾는 함수는 아래와 같이 작성하면 됩니다.

1
2
3
4
5
6
7
8
//query(root[i-1], root[j], 원소 범위 시작, 원소 범위 끝, k)
int query(Node *l, Node *r, int s, int e, int k){
    if(s == e) return s;
    int diff = r->l->v - l->l->v;
    int m = s + e >> 1;
    if(k <= diff) return query(l->l, r->l, s, m, k);
    else return query(l->r, r->r, m+1, e, k-diff);
}

BOJ11012 Egg

문제 링크

K번째 수와 비슷하지만, 약간 다른 관점에서 볼 수 있습니다.

점 추가 쿼리가 없는 2차원 구간 쿼리 문제에 PST를 활용할 수 있습니다.
세그먼트 트리를 x축, PST의 시간축을 y축으로 잡아서 2차원 평면을 표현해줄 수 있고, 쿼리 자체는 누적합 느낌으로 처리해주면 됩니다.