⛓ 자료구조
[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)
}
}