백준 - 구간 합 구하기 (2042)

import Foundation

let input = readLine()!.split(separator: " ").map { Int(String($0))! }
let (N, M, K) = (input[0], input[1], input[2])

// 세그먼트 트리 초기화
let height = ceil(log2(Double(N)))
let size = Int(pow(2.0, height + 1.0))
var tree = Array(repeating: 0, count: size + 1)

// 원본 데이터 넣기
setTree()

// 수 변경 또는 부분합 구하기
for _ in 0..<M + K {
    let input = readLine()!.split(separator: " ").map { Int(String($0))! }
    if input[0] == 1 {
        updateTree(input[1], input[2])
    } else {
        print(getSum(input[1], input[2]))
    }
}

// 세그먼트 트리 구축 함수
func setTree() {
    // 원본 데이터 넣기
    let startIndex = Int(pow(2.0, height))

    for i in startIndex..<startIndex + N {
        let input = Int(readLine()!)!
        tree[i] = input
    }
    
    // 부모 노드 값 설정 (부분합)
    var index = size
    while index != 1 {
        tree[index / 2] += tree[index]
        index -= 1
    }
}

// 세그먼트 트리 업데이트 함수
func updateTree(_ target: Int, _ value: Int) {
    var index = target + Int(pow(2.0, height)) - 1
    let diff = value - tree[index]
    while index > 0 {
        tree[index] += diff
        index /= 2
    }
}

// 세그먼트 트리 부분합 구하는 함수
func getSum(_ start: Int, _ end: Int) -> Int {
    var startIndex = start + Int(pow(2.0, height)) - 1
    var endIndex = end + Int(pow(2.0, height)) - 1
    var sum = 0
    while startIndex <= endIndex {
        if startIndex % 2 == 1 {
            sum += tree[startIndex]
        }
        if endIndex % 2 == 0 {
            sum += tree[endIndex]
        }
        startIndex = (startIndex + 1) / 2
        endIndex = (endIndex - 1) / 2
    }
    return sum
}

세그먼트 트리를 구현하는 기본적인 문제였다.

중간에 수의 변경이 빈번히 일어나기 때문에 효율적으로 부분 합을 구하기 위해 합 배열이 아닌 세그먼트 트리를 사용한다. 문제를 푼 김에 이번에 배운 세그먼트 트리 구현 방법을 정리하고자 한다.

세그먼트 트리의 구현 단계는 크게 트리 초기화, 트리 업데이트, 부분합 구하기 세 가지 단계로 나눌 수 있다.

트리 초기화

세그먼트 트리는 일차원 배열을 사용해서 구현한다. 이때 트리 배열의 크기는 2^k ≥ N을 만족하는 k의 최솟값을 구한 다음 2^(k+1)로 설정한다.

그리고 원본 데이터들을 트리의 리프 노드에 저장한다. 이때 시작 인덱스를 2^k 로 설정하고 차례대로 데이터를 삽입하면 된다.

트리 업데이트

세그먼트 트리는 이진 트리이기 때문에 인덱스의 값을 2로 나누면서 부모 노드를 찾아가며 업데이트 하면 된다.

부분 합 구하기

먼저 부분 합을 구할 시작과 끝 인덱스를 각각 트리의 리프 노드에 해당하는 인덱스로 바꿔주어야 한다. 이 인덱스는 주어진 인덱스 + 2^k - 1 로 설정한다.

그리고 리프 노드에서부터 올라가면서 부분 합을 구한다.

시작 인덱스가 홀수라면 이 노드는 오른쪽에 붙어 있는 노드이다. 따라서 부모 노드를 사용할 필요가 없기 때문에 이 노드의 값을 바로 더해준다.

끝 인덱스가 짝수라면 이 노드는 왼쪽에 붙어 있는 노드로, 마찬가지로 부모 노드를 사용할 필요가 없어 노드의 값을 바로 더해준다.

이 과정을 시작 인덱스와 끝 인덱스가 교차할 때까지 반복하면 부분 합을 구할 수 있다.

---

인덱스를 설정할 때 거듭제곱이랑 형 변환이 자주 일어나서 이 과정이 조금 헷갈리는 것 같다. 세그먼트 트리의 개념과 구현 방법을 잊지 않도록 나중에 꼭 복습을 해야겠다!