Heap.swift

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
    }
}