PriorityQueue

PriorityQueue.swift

struct PriorityQueue<Element>: CustomDebugStringConvertible {
  typealias PriorityFunc = (Element, Element) -> Bool
  private let priorityFunc: PriorityFunc
  private var elements: [Element]

  init(priorityFunction: @escaping PriorityFunc) {
    priorityFunc = priorityFunction
    elements = []
  }

  var count: Int { return elements.count }

  var isEmpty: Bool { return elements.isEmpty }

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

  mutating func pop() -> Element {
    assert(!isEmpty, "Queue is empty")
    elements.swapAt(0, count - 1)
    let result = elements.removeLast()
    siftDown(at: 0)
    return result
  }
  
  var peek: Element? { return elements.first }

  private mutating func siftUp(at index: Int) {
    var i = index
    while i != 0 {
      let p = parent(of: i)
      if !isValidPriority(p, i) {
        elements.swapAt(p, i)
        i = p
      } else {
        return
      }
    }
  }

  private mutating func siftDown(at index: Int) {
    guard index >= 0 && index <= (count - 1) / 2 else { return }
    var i = index
    let l = left(of: i)
    let r = right(of: i)
    if l < count && !isValidPriority(i, l) {
      i = l
    }
    if r < count && !isValidPriority(i, r) {
      i = r
    }
    if i != index {
      elements.swapAt(i, index)
      siftDown(at: i)
    }
  }

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

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

  private func right(of index: Int) -> Int {
    return index * 2 + 2
  }

  private func isValidPriority(_ lhs: Int, _ rhs: Int) -> Bool {
    return priorityFunc(elements[lhs], elements[rhs])
  }

  var debugDescription: String { return elements.debugDescription }
}