코틀린 코딩테스트 강좌, 세그먼트 트리

세그먼트 트리(Segment Tree)는 배열의 구간 합, 구간 최소값, 구간 최대값 등을 효과적으로 구할 수 있는 자료구조입니다. 이 자료구조는 특히 구간 쿼리와 업데이트가 빈번하게 이루어지는 상황에서 빠른 성능을 발휘합니다. 이 강좌에서는 세그먼트 트리의 기본적인 구조와 구현 방법, 그리고 실제 코딩테스트에서 활용할 수 있는 문제를 다루겠습니다.

세그먼트 트리란?

세그먼트 트리는 배열의 특정 구간에 대한 정보를 저장하기 위한 이진 트리 형태의 자료구조입니다. 각 노드는 특정 구간의 정보를 저장하고, 이를 통해 빠른 쿼리와 업데이트가 가능합니다. 주로 다음과 같은 두 가지 작업을 효율적으로 수행하는 데 사용됩니다.

  • 구간 쿼리(Query): 특정 구간의 합, 최소값, 최대값 등 계산
  • 업데이트(Update): 특정 위치의 값을 변경

세그먼트 트리 구조

세그먼트 트리는 완전 이진 트리로 구성되며, 일반적으로 다음과 같은 방식으로 배열을 구성합니다. 이를 위해, 세그먼트 트리의 각 노드는 부모와 자식 간의 관계를 통해 정의됩니다. 배열의 노드는 다음과 같은 규칙을 따릅니다:

  • 부모 노드는 인덱스 i의 왼쪽 자식이 2*i, 오른쪽 자식이 2*i + 1
  • 리프 노드는 주어진 배열의 각 원소와 상대하는 인덱스에 저장됩니다.

세그먼트 트리 구현

코틀린을 이용한 세그먼트 트리의 구현을 살펴보겠습니다. 먼저, 기본적인 세그먼트 트리 클래스를 정의하고, 초기화, 업데이트 및 쿼리 기능을 구현해 보겠습니다.

1. 세그먼트 트리 클래스 정의


class SegmentTree(private val array: IntArray) {
    private val size: Int = array.size
    private val tree: IntArray = IntArray(4 * size)
    
    init {
        build(0, 0, size - 1)
    }

    private fun build(node: Int, start: Int, end: Int) {
        if (start == end) {
            tree[node] = array[start]
        } else {
            val mid = (start + end) / 2
            build(2 * node + 1, start, mid)
            build(2 * node + 2, mid + 1, end)
            tree[node] = tree[2 * node + 1] + tree[2 * node + 2]
        }
    }

    fun update(index: Int, value: Int) {
        update(0, 0, size - 1, index, value)
    }

    private fun update(node: Int, start: Int, end: Int, index: Int, value: Int) {
        if (start == end) {
            array[index] = value
            tree[node] = value
        } else {
            val mid = (start + end) / 2
            if (index in start..mid) {
                update(2 * node + 1, start, mid, index, value)
            } else {
                update(2 * node + 2, mid + 1, end, index, value)
            }
            tree[node] = tree[2 * node + 1] + tree[2 * node + 2]
        }
    }
    
    fun query(L: Int, R: Int): Int {
        return query(0, 0, size - 1, L, R)
    }

    private fun query(node: Int, start: Int, end: Int, L: Int, R: Int): Int {
        if (R < start || end < L) return 0 // 쿼리 범위와 노드 범위가 겹치지 않음
        if (L <= start && end <= R) return tree[node] // 노드 범위가 쿼리 범위에 포함됨
        val mid = (start + end) / 2
        val leftSum = query(2 * node + 1, start, mid, L, R)
        val rightSum = query(2 * node + 2, mid + 1, end, L, R)
        return leftSum + rightSum
    }
}

2. 사용 예시

위에서 정의한 세그먼트 트리를 사용하여 구간의 합을 구하고 업데이트를 수행하는 예시를 살펴보겠습니다.


fun main() {
    val array = intArrayOf(1, 3, 5, 7, 9, 11)
    val segmentTree = SegmentTree(array)

    println("구간 [1, 3]의 합: ${segmentTree.query(1, 3)}") // 3 + 5 + 7 = 15
    segmentTree.update(1, 10) // array[1]를 10으로 업데이트
    println("업데이트 후 구간 [1, 3]의 합: ${segmentTree.query(1, 3)}") // 10 + 5 + 7 = 22
}

코딩테스트 문제

이제 실제 코딩테스트에서 자주 다뤄지는 문제를 함께 풀어보겠습니다. 아래 문제는 세그먼트 트리를 활용한 응용 문제입니다.

문제: 구간 합과 업데이트

주어진 배열의 특정 구간에 대한 합을 구하는 문제입니다. 추가로 특정 인덱스의 값을 업데이트할 수 있습니다.

문제 설명

- N개의 정수로 이루어진 배열이 주어집니다.
- 이전의 쿼리 결과로부터 Q번 만큼 다음의 작업을 수행합니다.
    1. "1 L R" 쿼리는 배열의 인덱스 L부터 R까지의 합을 요청합니다.
    2. "2 X Y" 쿼리는 인덱스 X의 값을 Y로 업데이트합니다.

입력 형식

- 첫 번째 줄에는 배열의 크기 N과 쿼리의 개수 Q가 주어집니다.
- 두 번째 줄에는 N개의 정수가 주어집니다.
- 그 다음 Q개의 줄에는 쿼리가 주어집니다.

출력 형식

- 각 "1 L R" 쿼리에 대해 구간의 합을 출력합니다.

입력 예시

5 3
1 2 3 4 5
1 1 3
2 1 10
1 1 3

출력 예시

6

문제 풀이

해당 문제를 해결하기 위해 세그먼트 트리를 사용하여 효율적으로 구간 합 쿼리와 업데이트 연산을 실행하도록 구현합니다. 입력으로부터 데이터를 받고 세그먼트 트리를 초기화한 후 쿼리를 처리하는 구조를 가집니다.

구현


fun main() {
    val (n, q) = readLine()!!.split(" ").map { it.toInt() }
    val array = readLine()!!.split(" ").map { it.toInt() }.toIntArray()
    val segmentTree = SegmentTree(array)

    repeat(q) {
        val query = readLine()!!.split(" ").map { it.toInt() }
        if (query[0] == 1) {
            // 구간 합 쿼리
            val sum = segmentTree.query(query[1] - 1, query[2] - 1)
            println(sum)
        } else if (query[0] == 2) {
            // 업데이트 쿼리
            segmentTree.update(query[1] - 1, query[2])
        }
    }
}

결론

이번 강좌에서는 세그먼트 트리의 기본 개념 및 구현 방법을 익혔습니다. 세그먼트 트리는 효율적인 구간 쿼리와 업데이트 구현에 매우 유용한 자료구조이며, 코딩테스트 문제에 자주 출제됩니다. 이번 예제를 통해 세그먼트 트리의 응용을 익혔기를 바라며, 이를 통해 향후 다양한 문제를 해결하는 데 도움이 되길 바랍니다.

코딩테스트 준비 시에는 이러한 다양한 자료구조와 알고리즘을 체계적으로 학습하여 응용 능력을 키우는 것이 중요합니다. 앞으로도 다양한 자료구조와 알고리즘을 통해 더 나은 문제 해결 능력을 갖추시기 바랍니다.