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

CS에서 statistic tree라고도 하는 세그먼트 트리는 간격 또는 세그먼트에 대한 정보를 저장하는 데 사용되는 트리 데이터 구조입니다.

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

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

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

 

1. 초기화

 

1차원 배열을 통해 구현한다.

* 0번째 idx는 사용하지 않는다.

 

원본 배열의 Size의 2배 크기의  배열을 Tree로 사용한다.

* 크기를 2배로 잡는 이유는 원본 배열로 세그먼트 트리의 리프 노드를 구성하게 되면,  

세그먼트 트리가 2진트리이기 때문에  (리프 노드 개수/ 2)면 모든 조상 노드들의 개수 보다 1만큼 크게 된다.

val originSumArr = arrayOf(5, 8, 4, 3, 10, 2, 1, 6)
lateinit var segmentTreeSum: Array<Int>

fun sumTreeInit(size: Int) {

    segmentTreeSum = Array(originSumArr.size * 2) { 0 }
    // leaf node 채우기
    for (i in 0 until originSumArr.size) {
        segmentTreeSum[i + segmentTreeSum.size / 2] = originSumArr[i]
    }
    // parent node 채우기
    for (i in (segmentTreeSum.size / 2) - 1 downTo 1)
        segmentTreeSum[i] = segmentTreeSum[i * 2] + segmentTreeSum[(i * 2) + 1]
    println("sumSegmentTree: ${segmentTreeSum.joinToString(" ")}")
}

2. 업데이트

값이 변경된 리프노드를 시작으로 2로 나눠가면서 조상노드를 갱신하면서, 새롭게 트리를 구성한다.

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

fun sumTreeUpdate(idx: Int, value: Int) {
    // 원본 변경
    originSumArr[idx] = value
    // 트리 변경
    var treeIdx = idx + (segmentTreeSum.size / 2)
    segmentTreeSum[treeIdx] = value
    // 리프노드부터 조상노드들에 차이값을 더 해줌
    while (treeIdx > 1) {
        treeIdx /= 2
        segmentTreeSum[treeIdx] = segmentTreeSum[treeIdx * 2] + segmentTreeSum[(treeIdx * 2) + 1]
    }
    println("sumSegmentTree: ${segmentTreeSum.joinToString(" ")}")
}

3. 쿼리

 

시작 ------->     <-------- 끝

 

시작 기준

"홀수" 이면 해당 노드만 독립적으로 결과값(ret)에 포함하고 왼쪽 조상노드로 이동을 위해 "+1"을 해주고 2로 나눈다.

"짝수" 이면 해당 노드가 조상노드에 포함되어 있으니까,  조상노드로 이동을 위해 2로 나눈다.

 

끝 기준

"홀수" 이면 해당 노드가 조상노드에 포함되어 있으니까,  조상노드로 이동을 위해 2로 나눈다.

"짝수" 이면 해당 노드만 독립적으로 결과값(ret)에 포함하고 왼쪽 조상노드로 이동을 위해 "-1"을 해주고 2로 나눈다.

 

fun sumTreeQuery(startIdx: Int, endIdx: Int): Int {

    var l = startIdx + (segmentTreeSum.size / 2)
    var r = endIdx + (segmentTreeSum.size / 2)

    var ret = 0
    while (l <= r) {
        if (l % 2 == 1) {
            ret += segmentTreeSum[l]
            l += 1
        }
        if (r % 2 == 0) {
            ret += segmentTreeSum[r]
            r -= 1
        }
        l /= 2
        r /= 2
    }
    return ret
}