Skip to content

Instantly share code, notes, and snippets.

@Nirma
Last active Dec 7, 2016
Embed
What would you like to do?
class Heap {
fileprivate var backingStore: [Int] = [Int]()
public var count: Int {
return backingStore.count
}
}
// MARK: - Indexing
private extension Heap {
func leftChild(of index: Int) -> Int? {
let leftIndex = (index * 2) + 1
return (count > leftIndex) && (leftIndex > 0) ? leftIndex : nil
}
func rightChild(of index: Int) -> Int? {
let rightIndex = (index * 2) + 2
return (count > rightIndex) && (rightIndex > 0) ? rightIndex : nil
}
func parent(of index: Int) -> Int? {
guard count > index, index > 0 else { return nil }
let parentIndex = (index - 1) / 2
return (count > parentIndex) && (parentIndex >= 0) ? parentIndex : nil
}
}
// MARK: - Ordering Heap
private extension Heap {
private func swap(_ indexA: Int, _ indexB: Int) {
(backingStore[indexA],backingStore[indexB]) = (backingStore[indexB], backingStore[indexA])
}
private func balanceSubHeap(at parentIndex: Int) -> Int? {
let parentValue = backingStore[parentIndex]
if let leftIndex = leftChild(of: parentIndex), let rightIndex = rightChild(of: parentIndex) {
let leftChild = backingStore[leftIndex]
let rightChild = backingStore[rightIndex]
if leftChild > max(parentValue, rightChild) {
swap(leftIndex, parentIndex)
return leftIndex
} else if rightChild > max(parentValue, leftChild) {
swap(rightIndex, parentIndex)
return rightIndex
}
} else if let leftIndex = leftChild(of: parentIndex) {
let childValue = backingStore[leftIndex]
if childValue > parentValue {
swap(leftIndex, parentIndex)
return leftIndex
}
}
return nil
}
func sortHeapDown() {
var current = 0
while let next = balanceSubHeap(at: current) {
current = next
}
}
func sortHeapUp() {
var currentIndex = count - 1
while let nextParent = parent(of: currentIndex) {
_ = balanceSubHeap(at: nextParent)
currentIndex = nextParent
}
}
}
// MARK: - Insertion & Deletion
extension Heap {
public func append(_ element: Int) {
backingStore.append(element)
sortHeapUp()
}
public func peak() -> Int? {
return backingStore.first
}
public func pop() -> Int? {
guard let head = backingStore.first else { return nil }
backingStore.removeFirst()
if let tail = backingStore.popLast() {
backingStore.insert(tail, at: 0)
sortHeapDown()
}
return head
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment