[Data Structure] 세그먼트 트리(Segment Tree)

세그먼트 트리?

• 여러 개의 데이터가 존재할 때 특정 범위의 합을 구하는데 사용하는 자료구조이다.
• 트리 종류 중에 하나로 이진트리의 형태이다.

0	1	2	3	4	5	6	7	8	9
5	6	9	3	8	2	5	1	4	7

• 위의 데이터에서 1부터 7까지의 합을 구한다 했을 때, 일반적으로 데이터를 하나씩 돌면서 더하면 O(N)의 시간이 걸린다. 하지만 세그먼트 트리를 사용하면 O(logN)의 시간에 구할 수 있다.

 

구현

 

세그먼트 트리 생성

먼저 전체 원소를 더한 값을 최상단 노드에 넣는다.

이때, 최상단 노드의 인덱스는 1을 부여한다.

트리를 구현하다보면 알겠지만, 양쪽 자식노드의 인덱스를 부여할 때 0부터 시작하는 것보다 1부터 시작하는 것이 편하다.

왼쪽자식노드 = 부모노드*2, 오른쪽자식노드 = 부모노드 * 2 + 1

 

이후, 두 번째 노드에는 0부터 4까지의 합을, 세 번째 노드에는 5부터 9까지의 합을 넣고 각각 왼쪽, 오른쪽 자식 노드로 연결한다.

이러한 과정을 반복하면 다음과 같이 전체 노드를 구할 수 있으며, 그 코드는 다음과 같다.

# start, end: 시작, 끝 인덱스
# 현재 인덱스(최상단 노드 인덱스인 1부터 시작)
int init(int start, int end, int nodeIndex) {
    if (start == end) {
        return tree[nodeIndex] = arr[start];
    }

    int mid = (start + end) / 2;

    return tree[nodeIndex] = init(start, mid, nodeIndex * 2) + init(mid + 1, end, nodeIndex * 2 + 1);
}

위 코드는 재귀적으로 수행한다.

참고로 주의할 점은 tree 크기를 미리 할당할 때 기존 배열 길이의 4배로 생성해주어야 index가 넘지 않는다.

 

구간 합 구하기

위에서 말했듯이 세그먼트 트리의 탐색은 O(logN)의 시간 복잡도를 가진다. 때문에 구간 합을 항상 O(logN)의 시간에 구할 수 있다.

# start, end: 시작, 끝 인덱스
# left, right: 구간 합 범위
# nodeIndex: 현재 노드 인덱스
int intervalSum(int start, int end, int nodeIndex, int left, int right) {
    // 범위 밖에 있는 경우
    if (left > end || right < start) {
        return 0;
    }

    // 범위 안에 있는 경우
    if (left <= start && right >= end) {
        return tree[nodeIndex];
    }

    // 그렇지 않으면 두 부분으로 나누어 합 구하기
    int mid = (start + end) / 2;
    return intervalSum(start, mid, nodeIndex * 2, left, right) + intervalSum(mid + 1, end, nodeIndex * 2 + 1, left, right);
}

위 코드는 재귀적으로 구현되었다.

코드를 한 번 살펴보면 시작과 끝 인덱스가 범위 안에 있을 때 해당 노드를 반환하며 범위가 아닌 경우(걸쳐있는 경우)에는 자식노드를 탐색한다.

left가 end보다 크거나 right가 start보다 작은 경우는 아예 범위 밖이기 때문에 0을 return 한다.

 

예를 들어 3부터 8까지의 합을 구한다고 하면 다음과 같이 전위순회하며 노드를 탐색하며 범위 내로 진입하면 노드를 반환한다.

최종적으로 다음과 같이 색칠된 노드를 반환한다.

 

구간 값 업데이트

// start, end: 구간 시작, 끝
// nodeIndex: 현재 노드 인덱스
// updateIndex: 업데이트 하려는 인덱스
// src: 업데이트 하려는 값
void update(int start, int end, int nodeIndex, int updateIndex, int src) {
    if (updateIndex < start || updateIndex > end) {
        return;
    }
    tree[nodeIndex] += src - arr[updateIndex];

    if (start == end) {
        return;
    }

    int mid = (start + end) / 2;
    update(start, mid, nodeIndex * 2, updateIndex, src);
    update(mid + 1, end, nodeIndex * 2 + 1, updateIndex, src);
}

위에서 비슷한 형태로 트리의 루트부터 방문하여 구간 합의 값을 계속 업데이트 해나간다.

tree[nodeIndex] += src - arr[updateIndex]; 에서 해당 인덱스의 기존 값을 제거하고 새로운 값을 더하는 식으로 업데이트 되는 방식이다.

여기서 주의할 점은 업데이트 수행 시, arr배열에서 해당 인덱스의 값도 변경해주어야 한다.

왜냐하면 arr의 값이 새로 반영되어야 다음 업데이트 때 새로 반영된 값을 기준으로 다시 업데이트 할 수 있기 때문이다.