Created
September 8, 2019 10:18
-
-
Save joanromano/58b1251c6de474bd45473cb70b36aa5d to your computer and use it in GitHub Desktop.
An array based implementation of a segment tree
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
extension Int { | |
fileprivate func nextPowerOf2() -> Int { | |
guard self != 0 else { | |
return 1 | |
} | |
return 1 << (bitWidth - (self - 1).leadingZeroBitCount) | |
} | |
} | |
struct SegmentTree<T: Comparable> { | |
private var array: [T] | |
private let priority: (T, T) -> T | |
private let inputMax: Int | |
init(_ elements: [T], priority: @escaping (T, T) -> T) { | |
precondition(!elements.isEmpty, "elements shouldn't be empty") | |
self.priority = priority | |
self.array = Array(repeating: elements.first!, count: elements.count.nextPowerOf2() * 2 - 1) | |
self.inputMax = elements.count - 1 | |
construct(input: elements, min: 0, max: elements.count-1, pos: 0) | |
} | |
private mutating func construct(input: [T], min: Int, max: Int, pos: Int) { | |
guard min != max else { | |
array[pos] = input[min] | |
return | |
} | |
let mid = min + (max - min) / 2 | |
construct(input: input, min: min, max: mid, pos: 2*pos+1) | |
construct(input: input, min: mid+1, max: max, pos: 2*pos+2) | |
array[pos] = priority(array[2*pos+1], array[2*pos+2]) | |
} | |
public func query(min: Int, max: Int) -> T? { | |
guard min >= 0, max <= inputMax else { return nil } | |
return query(qMin: min, qMax: max, min: 0, max: inputMax, pos: 0) | |
} | |
public mutating func update(index: Int, value: T) { | |
update(index: index, value: value, min: 0, max: inputMax, pos: 0) | |
} | |
private mutating func update(index: Int, value: T, min: Int, max: Int, pos: Int) { | |
if index < min || index > max { return } | |
if min == max { | |
array[pos] = value | |
return | |
} | |
let mid = min + (max - min) / 2 | |
update(index: index, value: value, min: min, max: mid, pos: 2*pos+1) | |
update(index: index, value: value, min: mid+1, max: max, pos: 2*pos+2) | |
array[pos] = priority(array[2*pos+1], array[2*pos+2]) | |
} | |
private func query(qMin: Int, qMax: Int, min: Int, max: Int, pos: Int) -> T? { | |
if qMin <= min, qMax >= max { // total overlap | |
return array[pos] | |
} | |
if qMin > max || qMax < min { // no overlap | |
return nil | |
} | |
// partial overlap | |
let mid = min + (max - min) / 2 | |
let left = query(qMin: qMin, qMax: qMax, min: min, max: mid, pos: 2*pos+1) | |
let right = query(qMin: qMin, qMax: qMax, min: mid+1, max: max, pos: 2*pos+2) | |
if let left = left, let right = right { | |
return priority(left,right) | |
} else if let left = left { | |
return left | |
} else if let right = right { | |
return right | |
} else { | |
return nil | |
} | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment