priority queue.swift

// MARK - Heap
public struct Heap<Element: Equatable> {
    var elements: [Element]
    let sort: (Element, Element) -> Bool
    
    public init(sort: @escaping (Element, Element) -> Bool, elements: [Element] = []) {
        self.sort = sort
        self.elements = elements
        
        // heapify
        if !elements.isEmpty {
            for i in stride(from: elements.count / 2 - 1, through: 0, by: -1) {
                sink(from: i)
            }
        }
    }

    public var isEmpty: Bool { elements.isEmpty }
    public var count: Int { elements.count }
    public func peek() -> Element? { elements.first }
    
    // append a new element
    public mutating func insert(_ element: Element) {
        // insert to the last position and bottom-up reheapify
        elements.append(element)
        swim(from: count - 1)
    }
    
    // remove and get the heap-top element
    public mutating func remove() -> Element? {
        guard !isEmpty else {
            return nil
        }
        elements.swapAt(0, count - 1)
        defer {
            sink(from: 0)
        }
        return elements.removeLast()
    }
    
    
    // Bottom-up reheapify implementation
    private mutating func swim(from index: Int) {
        var child = index
        var parent = parentIndex(ofChildAt: child)
        
        // swap child to upper position if sort(child, parent) is true
        while child > 0, sort(elements[child], elements[parent]) {
            elements.swapAt(child, parent)
            child = parent
            parent = parentIndex(ofChildAt: child)
        }
    }
    // Top-down reheapify implementation
    private mutating func sink(from index: Int) {
        var parent = index
        while true {
            let left = leftChildIndex(ofParentAt: parent)
            let right = rightChildIndex(ofParentAt: parent)
            
            var candidate = parent
            // choose one of (left, right)
            if (left < count && sort(elements[left], elements[candidate])) {
                candidate = left
            }
            if (right < count && sort(elements[right], elements[candidate])) {
                candidate = right
            }
            // no need to swap
            if candidate == parent {
                return
            }
            
            elements.swapAt(candidate, parent)
            parent = candidate
        }
    }
    
    // Index
    private func parentIndex(ofChildAt index: Int) -> Int {
        (index - 1) / 2
    }
    private func leftChildIndex(ofParentAt index: Int) -> Int {
        index * 2 + 1
    }
    private func rightChildIndex(ofParentAt index: Int) -> Int {
        index * 2 + 2
    }
}

public protocol Queue {
    associatedtype Element
    
    mutating func enqueue(_ element: Element) -> Bool
    mutating func dequeue() -> Element?
    
    var isEmpty: Bool { get }
    var peek: Element? { get }
}


public struct PriorityQueue<Element: Equatable>: Queue {
    private var heap: Heap<Element>
    
    init(sort: @escaping (Element, Element) -> Bool, elements: [Element] = []) {
        heap = Heap(sort: sort, elements: elements)
    }
    
    public mutating func enqueue(_ element: Element) -> Bool {
        heap.insert(element)
        return true
    }
    public mutating func dequeue() -> Element? {
        heap.remove()
    }
    
    public var isEmpty: Bool { heap.isEmpty }
    public var peek: Element? { heap.peek() }
}