Pasted

struct Heap<E: Comparable> {
    enum HeapType { case max, min }
 
    // an array representing the binary tree
    var tree = [E]()
 
    // indicating if this heap is a max heap or min heap
    let type: HeapType
 
    var count: Int {
        return tree.count
    }
 
    init(type: HeapType) {
        self.type = type
    }
 
    mutating func add(_ element: E) {
        tree.append(element)
        var child = tree.count - 1
        var parent = (child - 1) / 2
        while child > 0 && !satify(tree[parent], tree[child]) {
            tree.swapAt(parent, child)
            child = parent
            parent = (child - 1) / 2
        }
        assert(isHeap(0), "borken heap: \(tree)")
    }
 
    mutating func remove() -> E? {
        let rev = tree.first
        if tree.count > 0 {
            tree.swapAt(0, tree.count - 1)
            tree.removeLast()
            heapify(0)
        }
        assert(isHeap(0), "borken heap: \(tree)")
        return rev
    }
 
    func peek() -> E? {
        return tree.first
    }
 
    mutating private func heapify(_ rootIndex: Int) {
        let count = tree.count
        let leftChild = 2 * rootIndex + 1
        let rightChild = 2 * rootIndex + 2
 
        // no children
        if leftChild >= count { return }
 
        let targetChild = (rightChild >= count || satify(tree[leftChild], tree[rightChild])) ? leftChild : rightChild
        if (!satify(tree[rootIndex], tree[targetChild])) {
            tree.swapAt(rootIndex, targetChild)
            heapify(targetChild)
        }
    }
 
    private func satify(_ lhs: E, _ rhs: E) -> Bool {
        return (self.type == .min) ? (lhs <= rhs) : (lhs >= rhs)
    }
 
    // debugging helper methods
    private func isHeap(_ rootIndex: Int) -> Bool {
        let leftChild = 2 * rootIndex + 1
        let rightChild = 2 * rootIndex + 2
        var valid = true
        if leftChild < tree.count {
            valid = satify(tree[rootIndex], tree[leftChild]) && isHeap(leftChild)
        }
 
        if !valid { return false }
 
        if rightChild < tree.count {
            valid = satify(tree[rootIndex], tree[rightChild]) && isHeap(rightChild)
        }
        return valid
    }
}