[백준 2042번] [Kotlin] 구간 합 구하기

https://www.acmicpc.net/problem/2042

2042번: 구간 합 구하기

 

난이도 : 골드 1
태그 : 세그먼트 트리, 자료구

 

 

설명

세그먼트 트리를 활용한 구간 합 구하기 문제입니다.

 

 

세그먼트 트리란, 최상위 노드(루트)에는 모든 값들의 합,

그리고 범위를 이분할하여 해당 범위의 값을 저장하고,

그 자식 노드 역시 범위가 1이 될 때까지 반복적으로 분할합니다.

 

즉, 미리 구간합을 구해놓고 사용하는 방법으로써,

트리의 자료구조 상 중간의 값이 업데이트 된다 하여도 빠르게 값을 바꿀 수 있습니다.

 

 

소스코드

import java.util.*

lateinit var arr: Array<Long>
lateinit var tree: Array<Long>

fun main() = with(System.`in`.bufferedReader()) {
    // 수의 개수 m, 변경이 일어나는 횟수 m, 구간 합 구하는 횟수 k
    val (n, m, k) = readLine().split(" ").map { it.toInt() }

    // 첫 번째 값은 0으로 초기화
    arr = Array(n + 1) { if (it == 0) 0L else readLine().toLong() }
    tree = Array(n * 4 + 1) { 0L } // * 4 를 하면 모든 범위 커버 가능

    init(1, n, 1)

    val sb = StringBuilder()

    // 변경 및 합 출력
    repeat(m + k) {
        val st = StringTokenizer(readLine())
        // a, b, c
        // a - 1 : b를 c로 바꿔라
        // a - 2 : b 부터 c 까지 구간 합을 구하라
        val a = st.nextToken().toInt()
        val b = st.nextToken().toInt()

        // Update
        if (a == 1) {
            val c = st.nextToken().toLong()
            val diff = c - arr[b] // 차이 값 구하기
            arr[b] = c.toLong() // b 값을 c로 변경
            update(1, n, 1, b, diff) // 차이 값 만큼 갱신
        } else {
            val c = st.nextToken().toInt()
            // Sum
            sb.append(sum(1, n, 1, b, c)).append("\n")
        }
    }

    print(sb)
}


fun init(start: Int, end: Int, node: Int): Long {
    if (start == end) {
        tree[node] = arr[start]
        return tree[node]
    }

    val mid = (start + end) / 2
    tree[node] = init(start, mid, node * 2) + init(mid + 1, end, node * 2 + 1)
    return tree[node]
}


fun sum(start: Int, end: Int, node: Int, left: Int, right: Int): Long {
    // 범위 밖
    if (left > end || right < start) return 0
    // 범위 안
    if (left <= start && end <= right) return tree[node]
    // 그렇지 않다면 두 부분으로 나눠 합 구하기
    val mid = (start + end) / 2
    return sum(start, mid, node * 2, left, right) + sum(mid + 1, end, node * 2 + 1, left, right)
}

// index : 구간 합을 수정하고자 하느 노드
// dif : 수정할 값
fun update(start: Int, end: Int, node: Int, index: Int, diff: Long) {
    // 범위 밖
    if (index < start || index > end) return
    tree[node] += diff
    if (start == end) return
    val mid = (start + end) / 2
    update(start, mid, node * 2, index, diff)
    update(mid + 1, end, node * 2 + 1, index, diff)
}

 

init은 세그먼트 트리의 초기화,

sum은 특정 구간의 값을 구하는 메소드,

update는 특정 값을 바꾸는 메소드입니다.

 

 

후기

세그먼트 트리라는 처음보는 자료구조를 만났습니다.

도대체 이게 뭘까. 하면서 3시간동안 이것만 보고 있었더니 조금 개념은 잡힌 것 같네요.

조금 이해되고 나면, 나중에 알고리즘 카테고리에 세그먼트 트리를 주제로 포스팅을 해보고 싶네요.

지금은 일단 눈물부터 닦아야겠습니다. ㅠ.ㅠ