[자료구조] 세그먼트 트리 (Segment Tree)

세그먼트 트리

세그먼트 트리는 특정 구간의 최대값, 최소값, 합, 곱 등을 빠르고 간단하게 구할 수 있게 해주는 자료구조 입니다.

 

예를 들어
배열 arr = [1, 3, 2, 5, 7, 3, 5] 에서 3 ~ 7 번째 수 까지의 구간합을 구한다고 가정해봅시다.

 

가장 단순한 방법으로는 선형적으로 3번부터 8번까지의 수를 더하며 구할 수 있습니다.
하지만 이 방법은 $O(n)$으로 n = 1,000,000,000 처럼 크다면 상당히 부담되는 연산량입니다.

 

그래서 누적합을 미리 구해 구간합을 구하는 방법도 있습니다.
prefixSum = [0, 1, 4, 6, 11, 18, 21, 26] 을 구한 뒤 prefixSum[7] - prefixSum[2]을 계산하여 구간합을 얻을 수 있겠죠.
문제는 배열의 값을 수정하는 경우에 발생합니다. 배열의 값 하나를 바꾼다면 누적합 배열을 다시 새로 구해야 하는 일이 생기고 이때의 시간 복잡도 또한 $O(n)$으로, 배열의 값이 자주 수정된다면 부담스럽습니다.

 

 

이런 경우에 세그먼트 트리를 사용하는 것이 좋습니다.
세그먼트 트리는 구간합을 구하는데에 $O(\log{n})$, 값을 수정하는 데에도 $O(\log{n})$의 시간이 걸립니다.

 

세그먼트 트리가 어떻게 작동하는지 알아봅시다.

이진 트리의 형태로 구성되어 있고 각각의 노드는 구간의 값을 가지고 있습니다. 위의 예시를 따라 각 구간의 합을 가지고 있다고 가정해봅시다.
루트 노드는 1~7 번째 수들의 합을 저장하고, 그 자식노드는 각각 절반씩 구간을 나누어 그 구간의 합을 저장합니다.

매 자식마다 같은 방법으로 부모의 절반 구간 합을 저장하며 더이상 나눌 수 없을때 까지 자식을 가지게 됩니다.

앞선 예시 배열 arr = [1, 3, 2, 5, 7, 3, 5] 의 구간 합을 세그먼트 트리로 표현해보겠습니다.

init

 

 

get query

3~7 구간의 합을 구하고 싶다면 3~4 노드의 값 + 5~7 노드의 값 을 계산하여 구할 수 있습니다.

 

 

 

update

배열의 값을 수정하고 싶다면 루트부터 변경하고 싶은 값을 저장한 노드까지의 경로에 있는 노드들의 값만 수정해주면 됩니다.
예를 들어 4번째 값을 9로 수정해봅시다.

누적합을 사용할때보다 변경해야 할 값들이 훨씬 줄어들었음을 알 수 있습니다.

 

 

 

구현

세그먼트 트리가 어떻게 작동하는지 알아보았으니 이제 구현하는 방법을 알아봅시다.

마찬가지로 구간의 합을 나타내는 세그먼트 트리를 만들어 보겠습니다.

 

세그먼트 트리는 이진 트리이기 때문에 1차원 배열로 표현할 수 있습니다.

배열의 크기를 얼마로 잡아야 할까요?

 

입력 배열의 크기가 n 이라면 세그먼트 트리는 n개의 리프노드가 생기고, ?? 이진트리입니다.
따라서 $\lceil {\log{n}} \rceil +1$의 의 높이를 가진 트리가 생기게 됩니다.
높이 height를 가지는 트리의 노드 개수는 $2^{height}$, 1 << height 로 구할 수 있습니다.

    int height;
    int treeSize;
    vector<int> arr;
    vector<int> tree;

    arr.resize(n);
    height = ceil(log2(n))+1;
    treeSize = 1 << height;
    tree.resize(treeSize, 0);

n 보다 큰 중 가장 작은 2의 제곱수의 2배, 혹은 간단하게 n * 4 로 넉넉히 구할 수 있다.

 

 

이진 트리를 1차원 배열로 표현하는 경우 부모 - 자식 간의 관계는 다음과 같습니다.
루트가 1인 경우

• 부모 인덱스 : n
• 왼쪽 자식 인덱스 : n * 2
• 오른쪽 자식 인덱스 : n * 2 + 1

루트가 0인 경우

• 부모 인덱스 : n
• 왼쪽 자식 인덱스 : n * 2 + 1
• 오른쪽 자식 인덱스 : n * 2 + 2

 

루트의 인덱스가 1이냐, 0이냐 에 따라 자식노드에 접근하는 방법이 달라집니다.
우리는 루트가 1인 경우로 가정하고 구현해봅시다.

 

 

init

배열이 주어졌을 때, 우선 배열을 토대로 세그먼트 트리를 만들어야 합니다.

// start : 구간의 가장 처음, 배열 인덱스
// end : 구간의 가장 마지막, 배열 인덱스
// idx : 현재 노드의 인덱스, 트리 인덱스

int init(int idx, int start, int end)
{
    if (start == end)
        return tree[idx] = arr[start];
    
    int mid = (start + end) / 2;
    return tree[idx] = init(idx*2, start, mid) + init(idx*2 + 1, mid + 1, end);
}

//...

    init(1, 1, n);

• start : 구간의 가장 처음, 배열 인덱스를 따름
• end : 구간의 가장 마지막, 배열 인덱스를 따름
• idx : 현재 노드의 인덱스, 트리 인덱스를 따름

start == end 인 경우는 리프 노드에 해당합니다. 리프노드는 배열의 원소를 가집니다.

부모 노드의 값을 먼저 알 수 없기 때문에 자식 노드 먼저 구한 뒤 부모를 계산합니다.
왼쪽 자식노드는 구간의 앞부분 절반, 오른쪽 자식 노드는 구간의 뒷부분 절반을 담당하여 계산합니다.

트리를 만들기 위해서는 루트 인덱스, 루트의 구간을 파라미터로 넘겨 재귀적으로 트리를 만듭니다.

 

 

get query

원하는 구간에 대한 구간합을 계산해야 합니다.
합을 구해야하는 구간이 [ left, right ]로 주어진다면 트리를 순회하며 각 노드가 담당하는 구간과 left, right의 관계를 살펴봐야 합니다.

1. [ start, end ] 와 [ left, right ] 가 전혀 겹치지 않는 경우
2.  [ start, end ] 가 [ left, right ]에  완전히 포함 하는 경우, left <= start && end <= right
3. 그 외의 경우,  [ start, end ] 가 [ left, right ] 에 일부만 포함되는 경우

int getQuery(pair<int, int> range, int idx, int start, int end)
{
    if (end < range.first || range.second < start)
        return 0;
    
    if (range.first <= start && end <= range.second)
        return tree[idx];
    
    int mid = (start + end) / 2;
    return getQuery(range, idx*2, start, mid) + getQuery(range, idx*2 + 1, mid+1, end);
}

마찬가지고 구간합도 재귀함수 호출을 통해 구합니다.

 

1의 경우, 노드가 구해야 할 구간 합과 전혀 관련이 없으므로 0을 리턴합니다.

2의 경우, 노드가 구해야 할 구간 합의 일부를 가지고 있으므로 그 노드의 값을 리턴합니다.

3의 경우, 노드의 구간이 구해야 할 구간과 일부분만 겹치므로 완전히 겹치는 구간과 겹치지 않는 구간으로 나누기 위해 재귀적으로 절반씩 담당하는 자식 노드에서 다시 구간을 확인합니다. 

 

update

배열의 값이 변경될 경우 구간합의 값도 변경됩니다. 즉 변경된 수의 인덱스를 포함하는 구간의 노드들은 전부 업데이트 해주어야 합니다.
init에서 했던 방법처럼 리프노드의 값을 먼저 수정하고 그의 부모노드들을 차례로 수정하면 됩니다.

구간 합을 구하는 세그먼트 트리인 경우는 diff = 변경 후 - 변경 전 를 저장해 두고 배열 인덱스를 포함하는 노드들의 값에 diff 를 더해주어 해결할 수 있습니다.

 

int update(int num, int target, int idx, int start, int end)
{
    // if (start == end && target == start)
    //     return tree[idx] = num;

    if (target < start || end < target) // 반드시 이 조건 먼저 확인
        return tree[idx];

    if (start == end)
        return tree[idx] = num;

    int mid = (start + end) / 2;
    return tree[idx] = update(num, target, idx*2, start, mid) + update(num, target, idx*2 + 1, mid+1, end);
}

• num : 수정 후 값
• target : 변경될 배열의 인덱스

수정 해야할 인덱스가 현재 노드 범위에 해당하는지를 먼저 확인해야 합니다.
target 이 노드가 담당하는 구간에서 벗어난다면 해당 노드는 값이 변경되지 않습니다.

target이 노드의 범위 안에 해당하고, 그 노드가 리프노드라면 target == start == end 이고, 리프노드를 num 으로 변경하여 배열의 값을 변경할 수 있습니다.

부모 노드들은 재귀 함수 호출로 리프노드부터  루트까지 자식노드들의 합으로 변경합니다.