세그먼트 트리 (Segment Tree)

엄청나게 중요한 자료구조인 세그먼트 트리에 대하여 알아보고 가자.

세그먼트 트리의 목적은 구간에 대한 질의와 수정을 빠르게 하는 것이다.

 

아이디어는 다음과 같다.

구간 트리의 루트는 배열의 전체 구간인 [1, n]을 표현하며, 그 왼쪽 노드는 [1, mid], 오른쪽 노드는 [mid+1, n]을 포함하게 함으로서 재귀적으로 구간에 대한 정보를 이진 트리에 저장한다.

만약 구간에 대한 쿼리가 떨어진다면 그 쿼리의 구간에 해당하는 노드들을 몇개 잡아 그 그 부분에 대해서만 계산하면 된다.

 

세그먼트 트리 자체의 원리에 대한 자세한 내용은 다른 블로그나 종만북을 복습하도록 하고, 구현과 시간복잡도에 초점을 맞춰 보자.

*세그먼트 트리의 이론 및 구현 : https://www.acmicpc.net/blog/view/9

 

먼저 세그트리 클래스의 구현 코드는 다음과 같다. (구간 합 쿼리)

struct SegTree
{
    int n; vector<ll> t;

    SegTree(int n) : n(n), t(4*n+10) { init(1, 1, n); }

    void init(int node, int tl, int tr)
    {
        if(tl==tr) { t[node]=arr[tl]; return; }
        int mid=(tl+tr)/2;
        init(node*2, tl, mid);
        init(node*2+1, mid+1, tr);
        t[node]=t[node*2]+t[node*2+1];
    }

    ll query(int node, int tl, int tr, int l, int r)
    {
        if(tr<l || r<tl) return 0;
        if(l<=tl && tr<=r) return t[node];
        int mid=(tl+tr)/2;
        return query(node*2, tl, mid, l, r)+query(node*2+1, mid+1, tr, l, r);
    }

    void update(int node, int tl, int tr, int idx, ll val)
    {
        if(tl==tr) { t[node]=val; return; }
        int mid=(tl+tr)/2;
        if(idx<=mid) update(node*2, tl, mid, idx, val);
        else update(node*2+1, mid+1, tr, idx, val);
        t[node]=t[node*2]+t[node*2+1];
    }

    ll query(int l, int r) { return query(1, 1, n, l, r); }
    void update(int idx, int val) { update(1, 1, n, idx, val); }
};

 

1. init()

세그트리를 처음에 만드는 부분이다.

일단 생성자에서는 해당 배열을 받고, 세그트리의 크기를 4*n으로 결정한다. (4*n은 절대 터질 일 없는 매우 안전한 숫자이다.)

이제 직접 세그트리를 만들것인데, 해당 배열을 계속 들고 다니고, 지금 작업중인 node, 그 노드가 표현하는 구간 tl, tr을 인자로 갖고 다닌다.

리프 노드는 tl==tr 일 것이니 종료시키고, 나머지 경우에 대해서는 그 노드의 구간을 반으로 쪼개 자식 노드로 내려간다.

이때 시작 노드를 무조건 1로 설정했기 때문에 lc=node*2, rc=node*2+1이다.

시간 복잡도는 각 노드에 대해 O(1), 노드는 O(4N) 개 미만이니, O(N) 이다.

시간 복잡도 : O(N)

 

2. query()

쿼리를 해결하기 위해 필요한 정보는 지금 작업중인 node, 그 노드가 표현하고 싶은 구간 tl, tr, 우리가 알고 싶은 구간 l, r이다.

경우를 나누자.

1. [l, r]과 [tl, tr] 이 만나지 않는 경우 : 우리가 알고 싶어하는 구간과 완전히 관련 없으니, 병합 과정에서 무시될수 있는 값을 반환한다.

2. [l, r]안에 [tl, tr]이 완전히 포함되는 경우 : 우리가 알고싶어하는 구간에 완전히 포함된는 노드이니, 미리 계산한 트리의 값을 반환한다.

3. 나머지 경우에는 다시 반으로 쪼개 재귀적으로 호출한다.

최악의 경우에 mid, mid+1의 구간이 쿼리로 들어오면 바닥까지 2번 반복하므로 시간 복잡도는 O(lgN) 이다.

시간 복잡도 : O(lgN)

 

3. update()

노드 하나의 값을 바꾸고 그 노드에 의해 영향을 미친 나머지 노드들을 갱신해 준다.

마찬가지로 지금 작업중인 node, 그 노드가 표현하고 싶은 구간 tl, tr, 갱신된 위치 idx, 갱신된 값 val으로 설정한다.

전체 구간을 포함하는 루트 노드부터 내려가 idx를 포함하는 쪽으로 움직인다.

재귀 호출이 끝나면 그 노드의 자식 노드들은 갱신이 완료된 상태이니, t[node]를 갱신해 준다.

무조건 갱신된 위치까지만 갔다가 올라오니까 시간 복잡도는 당연히 O(lgN) 이다

시간 복잡도 : O(lgN)

 

전체 과정에서 merge 연산은 init(), update() 에만 사용되었다.

또한 클래스 마지막에 외부에서 query 와 update 를 쉽게 부를 수 있도록 인터페이스를 마련해 놓았다.

포인터를 이용해 구현하지말고 배열을 이용하는 것이 훨씬 빠르며, 값이 여러개일 때는 별도의 클래스를 사용하거나, 벡터를 여러개 만들어 이용한다.