⛓ 자료구조

[Swift] Segment Tree

akwlak 2022. 11. 1. 01:03
class SegmentTree {
    
    var arr: [Int]
    var tree: [Int]
    
    init(_ arr: [Int]) {
        self.arr = arr
        self.tree = [Int](repeating: 0, count: arr.count * 4)
        
        initTree(start: 0, end: arr.count - 1, node: 1)
    }
    
    @discardableResult
    private func initTree(start: Int, end: Int, node: Int) -> Int {
        if start == end {
            tree[node] = arr[start]
            return tree[node]
        }
        
        let mid = (start + end) / 2
        let leftNode = initTree(start: start, end: mid, node: node * 2)
        let rightNode = initTree(start: mid + 1, end: end, node: node * 2 + 1)
        
        tree[node] = leftNode + rightNode
        
        return tree[node]
    }
    
    func sum(start: Int, end: Int, node: Int, left: Int, right: Int) -> Int {
        if left > end || right < start { return 0 }
        
        if left <= start && end <= right { return tree[node] }
        
        let mid = (start + end) / 2
        let leftSum = sum(start: start, end: mid, node: node * 2, left: left, right: right)
        let rightSum = sum(start: mid + 1, end: end, node: node * 2 + 1, left: left, right: right)
        
        return leftSum + rightSum
    }
    
    func update(start: Int, end: Int, node: Int, index: Int, diff: Int) {
        if index < start || index > end { return }
        
        tree[node] += diff
        
        if start == end { return }
        
        let mid = (start + end) / 2
        update(start: start, end: mid, node: node * 2, index: index, diff: diff)
        update(start: mid + 1, end: end, node: node * 2 + 1, index: index, diff: diff)
    }
}