55 lines
1.3 KiB
Swift
55 lines
1.3 KiB
Swift
struct Heap<T> {
|
|
var comparator: ((_ l: T,_ r: T) -> Bool)
|
|
var heap: [T] = []
|
|
var isEmpty: Bool { return heap.isEmpty }
|
|
|
|
private mutating func bubbleUp(idx: Int) {
|
|
let parent = (idx - 1) / 2
|
|
if idx <= 0 {
|
|
return
|
|
}
|
|
if comparator(heap[idx], heap[parent]) {
|
|
heap.swapAt(parent, idx)
|
|
bubbleUp(idx: parent)
|
|
}
|
|
}
|
|
|
|
private mutating func heapify(_ idx: Int) {
|
|
let left = idx * 2 + 1
|
|
let right = idx * 2 + 2
|
|
var comp = idx
|
|
|
|
if heap.count > left && comparator(heap[left], heap[comp]) {
|
|
comp = left
|
|
}
|
|
if heap.count > right && comparator(heap[right], heap[comp]) {
|
|
comp = right
|
|
}
|
|
if comp != idx {
|
|
heap.swapAt(comp, idx)
|
|
heapify(comp)
|
|
}
|
|
}
|
|
|
|
mutating func insert(_ item: T) {
|
|
heap.append(item)
|
|
bubbleUp(idx: heap.count-1)
|
|
}
|
|
|
|
mutating func pop() -> T? {
|
|
let item: T? = heap.first
|
|
if heap.count == 0 {
|
|
return nil
|
|
}
|
|
if heap.count > 1 {
|
|
heap[0] = heap[heap.count-1]
|
|
heap.removeLast()
|
|
heapify(0)
|
|
} else if heap.count == 1{
|
|
heap.removeLast()
|
|
}
|
|
return item
|
|
}
|
|
}
|
|
|