struct Heap { 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 } }