자료구조

Persistent Segment Tree

lovinix

2019년 11월 8일 • 7 min read

다음과 같은 문제를 생각해보자.

크기가 N인 배열에서 주어지는 [l,r] 범위들의 합을 각각 구하여라.

어디서 본 것같지 않은가? 그렇다. hld 글을 시작했던 그 문제이다.
저 글에서 우리는 이 문제에 update를 얹고 배열을 트리형태로 바꾸었을 때 어떻게 처리해야할 지를 다루었는데 이번에는 문제를 약간 바꿔보겠다.

크기가 N인 배열에서 주어지는 [l,r] 범위의 k번째로 작은 값을 구하여라.

[l,r] 범위가 없이 k번째 값만 찾는다고 생각해보면 배열 내의 숫자값들을 범위로 하는 세그먼트트리를 사용하여 kth element를 찾아낼 수 있다. 이 때 우리가 사용했던 세그먼트 트리를 배열의 [0,i]번째 인덱스까지 가지고있는 세그먼트트리로 만든 다음, [0,0], [0,1], [0,2], ... [0,r]까지 모두 가지고 있다면 [l,r] 범위의 k번째 값도 구할 수 있지 않을까?

현재 우리가 보고있는 숫자의 범위를 [s,e]라 하고 배열의 r번째 인덱스까지 저장하고있는 세그먼트트리에서 동일한 구간에 있는 숫자의 개수를 v_r, l-1번째 인덱스까지 저장하고 있는 세그먼트트리의 [s,e] 구간에 있는 숫자의 개수를 v_l-1라 하면 우리가 원하는 [l,r] 구간 내에 존재하는 숫자의 개수는 v_r - v_l-1가 될테니 [0,l-1] 세그먼트트리와 [0,r] 세그먼트트리를 가지고있다면 [l,r] 범위의 k번째 원소를 충분히 찾아내면서 O(logN)의 시간도 보장함을 알 수 있다.

하지만 이 방법의 치명적인 문제점은 바로 공간복잡도가 터져버린다는 것이다. 숫자의 개수를 N, 숫자의 범위를 M이라 하면 동적세그먼트트리를 쓴다해도 세그먼트트리 하나를 구축하는데 NlogM정도의 공간이 필요한데 이게 N개 있으니 공간복잡도가 N²logM이 나와버리고 이러면 쿼리당 시간복잡도를 줄인 의미가 없어진다.

공간복잡도까지 줄일 방법이 없을까?

간략하게 보기위해 다음과같은 조건으로 세그먼트 트리를 만들어보자.

배열 : {1,3} , 숫자의 범위 [0,3]

먼저 0번째 원소(1)만 들어가있는 세그먼트트리를 그려보면 아래와 같다.
seg1
원의 위 또는 아래에 써있는 범위는 해당 노드에서의 숫자의 범위이고, 원 안의 숫자는 해당 구간에서의 숫자의 개수이다. 기본적으로 동적세그먼트트리로 짤테니 비어있는 부분은 아직 만들어지지 않은 노드라고 생각하자.

다음으로 이 세그먼트트리에 1번째 원소(3)까지 넣은 세그먼트트리를 그려보면 아래와 같다.
seg2

둘을 비교해서 변화가 있던 부분만을 칠해보면 아래와같은데
seg3

M이 숫자의 범위라고했으니 정확히 logM개의 노드에만 변화가 있었다. 그렇다면 나머지 노드는 전부 저장하는 것이 아니라 이전에 만들었던 노드를 가리키게하면 logM의 공간만 추가로 사용하면서도 새로운 세그먼트트리를 만들 수 있다.

seg4--1-

위와같은 형태로 저장해준다면 배열 인덱스 하나당 logM개의 노드만 생기며 전체 원소의 개수가 N개이므로 전체 공간복잡도는 NlogM이 될 것이다.

실제 구현은 위에서 정의한 [l,r]구간의 K번째 수 찾기를 한다고 생각하고 보도록 하자.

먼저 Segment tree의 노드가 될 구조체를 정의해주고, N+1개의 노드를 선언해준다.

struct Pst
{
    int l, r, v;
    Pst() : l(0), r(0), v(0) {}
};

vector<Pst> pst;

int main()
{
    pst.resize(N + 1);
    ...
}

노드에서 l,r은 각각 자신의 왼쪽자식, 오른쪽자식의 인덱스번호이며 선언한 N+1개의 원소 중 첫 1개 (0번 index)는 null의 역할을 하는 노드이고 나머지 N개는 각각 [1,i]를 (아래로는 1-based index를 쓰도록 하겠다) 담당하는 세그먼트트리의 루트이다.

다음으로 update이다. i번째 인덱스를 추가하는 세그먼트트리는 i-1번째 인덱스까지 추가한 세그먼트트리에 원소 하나를 추가하는 형태일테니 기본적으로 i-1번째세그먼트트리(i1)와 i번째 세그먼트트리(i2)를 동시에 타고 내려간다고 생각하자.

/*
    i1 : 이전에 참조할 노드의 번호
    i2 : 현재 보고있는 노드의 번호
    v : 세그먼트트리에 추가하려는 숫자
    s,e : 현재 노드에서 보고있는 숫자의 범위 [s,e]
*/
void update(int i1, int i2, int v, int s, int e)
{
    pst[i2].v = pst[i1].v + 1; // 현재 노드에 숫자 한 개 추가
    if(s==e) return;
    int m = s + e >> 1;
	if(v<=m)
	{
        /* 왼쪽 자식으로 가야하는데 왼쪽 자식이 아직 안만들어졌거나
           이전의 위치와 같은 경우 갱신해야하므로 노드를 새로 만들어준다 */
		if(pst[i2].l==0 || pst[i2].l==pst[i1].l)
		{
			pst[i2].l = pst.size();
			pst.emplace_back(Pst());
		}
        /* 오른쪽자식은 갱신될 일이 없으므로
           이전 세그먼트트리의 오른쪽 자식의 위치를 그대로 가져가자 */
		if (pst[i2].r == 0)
			pst[i2].r = pst[i1].r;
		update(pst[i1].l, pst[i2].l, v, s, m);
	}
	else
	{
		if(pst[i2].r==0 || pst[i2].r==pst[i1].r)
		{
			pst[i2].r = pst.size();
			pst.emplace_back(Pst());
		}
		if (pst[i2].l == 0)
			pst[i2].l = pst[i1].l;
		update(pst[i1].r, pst[i2].r, v, m + 1, e);
	}
}

동적세그먼트트리를 알고있다면 형태가 크게 다르지 않다는 것이 바로 보일것이다. 주석에 달려있는 것처럼 갱신을 해야하는데 노드가 없거나, 이전의 위치와 같은 경우에만 새로 만들고, 갱신하지 않는 쪽 노드는 이전의 위치를 그대로 저장함에 주의하자.

int query(int i1, int i2, int k, int s, int e)
{
	if (s == e) 
		return s;
	int m = s + e >> 1;
	int cntLeft = pst[pst[i2].l].v - pst[pst[i1].l].v;
    if(cntLeft>=k) return query(pst[i1].l, pst[i2].l, k, s, m);
    else return query(pst[i1].r, pst[i2].r, k-cntLeft, m+1, e);
}

[s,m] 구간에 있는 숫자의 개수는 pst[pst[i2].l].v - pst[pst[i1].l].v 와 같음을 이해한다면 기본 세그먼트트리에서 kth element를 찾는 과정과 완전히 동일하다.

Sign up for more like this.