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

1. 세그먼트 트리란?

 

간격 또는 세그먼트에 대한 정보를 저장하는데 사용되는 트리 데이터 구조이다.

주어진 포인트를 포함하는 저장된 세그먼트에 대해서 쿼리할 수 있다.

 

N개의 원소에 대하여 구성된 세그먼트 트리는 트리를 생성하는 비용은 O(N)이다.

원소를 Update하는데 O(Log N)의 시간 복잡도가 필요하고, 쿼리 또한 O(Log N)의 시간 복잡도가 요구된다.

 

세그먼트 트리는 누적합, 구간합, 구간의 최솟값을 구하는데 사용할 수 있다.

 

1.1. 생성

1차원 배열을 통해 구현하며 0번 Index는 사용하지 않는다.

원본 배열 Size의 2배 크기로 구성한다.

크기를 2배로 잡는 것은, 원본 배열로 세그먼트 트리의 리프 노드를 구성하게 되면, 세그먼트 트리가 이진트리이므로 (리프노드 개수 / 2)를 하면 모든 루트노드들의 개수보다 1만큼 크게 된다.

fun init(node: Int, start: Int, end: Int): Long {
    return if (start == end) {
        tree[node] = arr[start]
        tree[node]
    } else {
        val mid = (start + end) / 2
        tree[node] = init(node * 2, start, mid) + init(node * 2 + 1, mid + 1, end)
        tree[node]
    }
}

 

 

1.2. 갱신

index번째 수를 X로 변경하는 경우, index번째를 포함하는 노드에 들어있는 합만 변경하면 된다.

원래 index번째 수가 arr[index]였고, 바뀐 수가 X라면, 합은 X - arr[index]만큼 변한다.

 

수 변경은 다음과 같이 2가지 경우가 존재한다.

1. start, end에 index가 포함된 경우
2. start, end에 index가 포함되지 않은 경우

1번의 경우 재귀 호출, 2번은 그 노드의 모든 자식도 index번째를 포함하지 않으니 중단한다.

이때 이진트리 기반의 세그먼트 트리는 값의 변경에 O(Log N)의 시간복잡도가 요구된다.

 

index번째를 포함하는 모든 노드의 합에 diff를 더해서 수를 변경한다.

fun update(node: Int, index: Int, start: Int, end: Int, diff: Long) {
    if (index < start || index > end) return

    tree[node] += diff

    if (start != end) {
        val mid = (start + end) / 2
        update(node * 2, index, mid, end, diff)
        update(node * 2 + 1, index, mid + 1, end, diff)
    }
}

 

 

1.3. 구간합

구간 left, right가 주어졌을 때, 합을 구하려면 트리를 루트부터 순회하며 각 노드에 저장된 구간의 정보와 left, right의 관계를 살펴봐야한다. node에 저장된 구간이 start, end이고 합을 구해야하는 구간이 left, right라면 네 가지 경우가 나올 수 있다.

 

1. left, right와 start, end가 겹치치 않는 경우
2. left, right와 start, end를 완전히 포함하는 경우
3. start, end가 left, right를 완전히 포함하는 경우
4. left, right와 start, end가 겹쳐져 있는 경우

1번은 if (left > end || right < start)로 구현이 가능하다.

2번은 if (left <= start & end <= right)로 구현이 가능하다.

3번과 4번은 왼쪽 자식과 오른쪽 자식을 루트로 하는 트리에서부터 다시 탐색해야한다.

그렇기에 값을 바꾸어 재귀 호출한다.

fun sum(node: Int, start: Int, end: Int, left: Int, right: Int): Long {
    return when {
        left > end || right < start -> 0
        left <= start && end <= right -> tree[node]
        else -> {
            val mid = (start + end / 2)
            sum(node * 2, start, mid, left, right) + sum(node * 2 + 1, mid + 1, end, left, right)
        }
    }
}

 

 

전체 코드

class SegmentTree(private val arr: LongArray) {
    private val tree = LongArray(arr.size * 2) { 0L }

    fun init(node: Int, start: Int, end: Int): Long {
        return if (start == end) {
            tree[node] = arr[start]
            tree[node]
        } else {
            val mid = (start + end) / 2
            tree[node] = init(node * 2, start, mid) + init(node * 2 + 1, mid + 1, end)
            tree[node]
        }
    }

    fun update(node: Int, index: Int, start: Int, end: Int, diff: Long) {
        if (index < start || index > end) return

        tree[node] += diff

        if (start != end) {
            val mid = (start + end) / 2
            update(node * 2, index, mid, end, diff)
            update(node * 2 + 1, index, mid + 1, end, diff)
        }
    }

    fun sum(node: Int, start: Int, end: Int, left: Int, right: Int): Long {
        return when {
            left > end || right < start -> 0
            left <= start && end <= right -> tree[node]
            else -> {
                val mid = (start + end / 2)
                sum(node * 2, start, mid, left, right) + sum(node * 2 + 1, mid + 1, end, left, right)
            }
        }
    }
}

 

https://book.acmicpc.net/ds/segment-tree

 

https://codingnojam.tistory.com/49

[Algorithm] 세그먼트 트리(Segment Tree)를 Java로 구현해보자! !(with BOJ-2042 Java로 문제풀이)

https://hongjw1938.tistory.com/20

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

 

https://hik-coding.tistory.com/272?category=939138

세그먼트 트리(Segment Tree) feat. kotlin