Heap Swift

Heap.swift

struct Heap<Element> {
    var elements: [Element]
    let priorityFunction: (Element, Element) -> Bool

    init(elements: [Element] = [], priorityFunction: @escaping (Element, Element) -> Bool) {
        self.priorityFunction = priorityFunction
        self.elements = []

        for element in elements {
            enqueue(element)
        }
    }

    var isEmpty: Bool {
        return elements.isEmpty
    }

    var count: Int {
        return elements.count
    }

    func peek() -> Element? {
        return elements.first
    }

    func isRoot(_ index: Int) -> Bool {
        return index == 0
    }

    func leftChildIndex(of index: Int) -> Int {
        return 2 * index + 1
    }

    func rightChildIndex(of index: Int) -> Int {
        return leftChildIndex(of: index) + 1
    }

    func parentIndex(of index: Int) -> Int {
        return (index - 1) / 2
    }

    func isHigherPriority(at firstIndex: Int, than secondIndex: Int) -> Bool {
        return priorityFunction(elements[firstIndex], elements[secondIndex])
    }

    func highestPriorityIndex(of parentIndex: Int, and childIndex: Int) -> Int {
        guard childIndex < count && isHigherPriority(at: childIndex, than: parentIndex) else { return parentIndex }

        return childIndex
    }

    func highestPriorityIndex(for parent: Int) -> Int {
        return highestPriorityIndex(of: highestPriorityIndex(of: parent, and: leftChildIndex(of: parent)), and: rightChildIndex(of: parent))
    }

    mutating func swapElement(at firstIndex: Int, with secondIndex: Int) {
        guard firstIndex != secondIndex, 0..<count ~= firstIndex, 0..<count ~= secondIndex else {
            return
        }

        elements.swapAt(firstIndex, secondIndex)
    }

    mutating func enqueue(_ element: Element) {
        elements.append(element)
        siftUp(elementAt: count - 1)
    }

    mutating func siftUp(elementAt index: Int) {
        let parent = parentIndex(of: index)

        guard !isRoot(index), isHigherPriority(at: index, than: parent) else {
            return
        }

        swapElement(at: index, with: parent)
        siftUp(elementAt: parent)
    }

    mutating func dequeue() -> Element? {
        swapElement(at: 0, with: count - 1)
        let element = elements.popLast()

        siftDown(elementAt: 0)

        return element
    }

    mutating func siftDown(elementAt index: Int) {
        let child = highestPriorityIndex(for: index)
        if index == child {
            return
        }

        swapElement(at: index, with: child)
        siftDown(elementAt: child)
    }
}