BOJ-2042 구간 합 구하기 Swift

문제

어떤 N개의 수가 주어져 있다. 그런데 중간에 수의 변경이 빈번히 일어나고 그 중간에 어떤 부분의 합을 구하려 한다. 만약에 1,2,3,4,5 라는 수가 있고, 3번째 수를 6으로 바꾸고 2번째부터 5번째까지 합을 구하라고 한다면 17을 출력하면 되는 것이다. 그리고 그 상태에서 다섯 번째 수를 2로 바꾸고 3번째부터 5번째까지 합을 구하라고 한다면 12가 될 것이다.

입력

첫째 줄에 수의 개수 N(1 ≤ N ≤ 1,000,000)과 M(1 ≤ M ≤ 10,000), K(1 ≤ K ≤ 10,000) 가 주어진다. M은 수의 변경이 일어나는 횟수이고, K는 구간의 합을 구하는 횟수이다. 그리고 둘째 줄부터 N+1번째 줄까지 N개의 수가 주어진다. 그리고 N+2번째 줄부터 N+M+K+1번째 줄까지 세 개의 정수 a, b, c가 주어지는데, a가 1인 경우 b(1 ≤ b ≤ N)번째 수를 c로 바꾸고 a가 2인 경우에는 b(1 ≤ b ≤ N)번째 수부터 c(b ≤ c ≤ N)번째 수까지의 합을 구하여 출력하면 된다.

입력으로 주어지는 모든 수는 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

출력

첫째 줄부터 K줄에 걸쳐 구한 구간의 합을 출력한다. 단, 정답은 -2^63보다 크거나 같고, 2^63-1보다 작거나 같은 정수이다.

 

내가 푼 풀이

접근방법: 세그먼트 트리

구간합트리인 세그먼트 트리를 구현후, 해당 동작들을 구현하였다.

입력 및 출력은 Int64의 최소최대값이 된다. 그래서 더해도 이 범위를 넘지않으므로 따로 오버플로우에 대한 대응은 하지않았다.

다만 Int로 선언시 본인의 컴퓨터에따라 다른 최댓값과 최솟값을 가질수 있으므로 안전하게 Int64로 사용해도 된다.(본인은 Int사용)

원소의 값을 변경할때, 세그먼트 트리도 변경되지만 주어진 N개의 배열도 값을 수정해줘야 다음 값변경에도 적용이 되므로 꼭 수정해줘야한다.

 

import Foundation

// 입력받기
let input = readLine()!.split(separator: " ").map{Int(String($0))!}
let N = input[0], M = input[1], K = input[2]
var arr = [0]
for i in 0..<N {
    arr.append(Int(readLine()!)!)

}
// 세그먼트 트리 크기지정
var tree = Array(repeating: 0, count: arr.count * 4)

// 세그먼트 트리 생성함수
func sgInit(_ start: Int,_ end: Int,_ index: Int) -> Int{
    if start == end {
        tree[index] = arr[start]
        return tree[index]
    }
    let mid = (start + end)/2
    tree[index] = sgInit(start, mid, index*2) &+ sgInit(mid+1, end, index*2+1)
    
    return tree[index]
}
// 세그먼트트리를 이용한 구간합 구하기
func sgSum(_ start: Int,_ end: Int,_ index: Int,_ left: Int,_ right: Int) -> Int {
    if left > end || right < start { return 0 }
    if start >= left && end <= right {
        return tree[index]
    }
    let mid = (start+end)/2
    return sgSum(start, mid, index*2, left, right) &+ sgSum(mid+1, end, index*2+1, left, right)
}
// 원소값 변경으로 세그먼트트리의 값 수정
func sgUpdate(_ start: Int,_ end: Int,_ index: Int,_ node: Int,_ value: Int) {
    if start > node || end < node { return }
    tree[index] += value
    if start == end { return }
    let mid = (start+end)/2
    sgUpdate(start, mid, index*2, node, value)
    sgUpdate(mid+1, end, index*2+1, node, value)

}
// 세그먼트 트리 만들기
sgInit(1, arr.count-1, 1)

// 동작 수행
for i in 0..<M+K {
    let op = readLine()!.split(separator: " ").map{Int(String($0))!}
    if op[0] == 1 {
        // 세그먼트 트리의 값 수정해주고, 배열값도 수정해준다.
        sgUpdate(1, arr.count-1, 1, op[1], op[2] - arr[op[1]])
        arr[op[1]] = op[2]
    } else if op[0] == 2 {
        print(sgSum(1, arr.count-1, 1, op[1], op[2]))
    }
}