세그먼트 트리(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])
}
}
}
결론
이번 강좌에서는 세그먼트 트리의 기본 개념 및 구현 방법을 익혔습니다. 세그먼트 트리는 효율적인 구간 쿼리와 업데이트 구현에 매우 유용한 자료구조이며, 코딩테스트 문제에 자주 출제됩니다. 이번 예제를 통해 세그먼트 트리의 응용을 익혔기를 바라며, 이를 통해 향후 다양한 문제를 해결하는 데 도움이 되길 바랍니다.
코딩테스트 준비 시에는 이러한 다양한 자료구조와 알고리즘을 체계적으로 학습하여 응용 능력을 키우는 것이 중요합니다. 앞으로도 다양한 자료구조와 알고리즘을 통해 더 나은 문제 해결 능력을 갖추시기 바랍니다.