Skip to content

Instantly share code, notes, and snippets.

@rockbruno
Created October 8, 2020 17:19
Show Gist options
  • Save rockbruno/9eac854c74d2998d2ad1b0b24670e1c6 to your computer and use it in GitHub Desktop.
Save rockbruno/9eac854c74d2998d2ad1b0b24670e1c6 to your computer and use it in GitHub Desktop.
PriorityQueue
final class Heap<T> {
typealias Comparator = (T,T) -> Bool
var elements: [T]
let priority: Comparator
init(elements: [T], priority: @escaping Comparator) {
self.priority = priority
self.elements = elements
if elements.isEmpty == false {
for i in stride(from: (count / 2) - 1, to: -1, by: -1) {
siftDown(i)
}
}
}
var isEmpty: Bool {
return elements.isEmpty
}
var count: Int {
return elements.count
}
var first: T? {
return elements.first
}
func leftChildIndex(of index: Int) -> Int {
return (2 * index) + 1
}
func rightChild(of index: Int) -> Int {
return (2 * index) + 2
}
func parentIndex(of index: Int) -> Int {
return (index - 1) / 2
}
func isHigherPriority(_ a: Int, _ b: Int) -> Bool {
return priority(elements[a], elements[b])
}
func highestPriorityIndex(of index: Int) -> Int {
let left = highestPriorityIndex(of: index, and: leftChildIndex(of: index))
let right = highestPriorityIndex(of: index, and: rightChild(of: index))
return highestPriorityIndex(of: left, and: right)
}
func highestPriorityIndex(of parent: Int, and child: Int) -> Int {
guard child < count else {
return parent
}
guard isHigherPriority(child, parent) else {
return parent
}
return child
}
func enqueue(_ element: T) {
elements.append(element)
siftUp(count - 1)
}
func siftUp(_ i: Int) {
let parent = parentIndex(of: i)
guard parent >= 0 else {
return
}
guard isHigherPriority(i, parent) else {
return
}
swap(i, parent)
siftUp(parent)
}
func dequeue() -> T? {
guard count > 0 else {
return nil
}
swap(0, count - 1)
let element = elements.popLast()
siftDown(0)
return element
}
fileprivate func swap(_ i: Int, _ j: Int) {
(elements[i], elements[j]) = (elements[j], elements[i])
}
func siftDown(_ i: Int) {
let indexToSwap = highestPriorityIndex(of: i)
guard indexToSwap != i else {
return
}
swap(indexToSwap, i)
siftDown(indexToSwap)
}
}
Heap<String>(elements: []) { (a, b) -> Bool in
a.count < b.count
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment