Minimal Priority Queue Ruby Implementation

pqueue.rb

# PriorityQueue: binary-heap implementation (min-heap by default).
# API:
#   push(value, priority)
#   pop            # => value with smallest priority (or nil if empty)
#   peek           # => next value (or nil)
#   peek_priority  # => next priority (or nil)
#   size, empty?

class PriorityQueue
  def initialize(&comparator)
    @heap = []
    # comparator compares priorities: return true if a is "before" b.
    @cmp = comparator || ->(a, b) { a < b }  # min-heap by default
  end

  def push(value, priority)
    @heap << [priority, value]
    sift_up(@heap.size - 1)
    self
  end

  def pop
    return nil if @heap.empty?
    swap(0, @heap.size - 1)
    pr, val = @heap.pop
    sift_down(0) unless @heap.empty?
    val
  end

  def peek
    @heap.empty? ? nil : @heap[0][1]
  end

  def peek_priority
    @heap.empty? ? nil : @heap[0][0]
  end

  def size
    @heap.size
  end

  def empty?
    @heap.empty?
  end

  private

  def sift_up(i)
    while i > 0
      p = (i - 1) / 2
      break if @cmp.call(@heap[p][0], @heap[i][0])
      swap(i, p)
      i = p
    end
  end

  def sift_down(i)
    n = @heap.size
    loop do
      l = 2 * i + 1
      r = 2 * i + 2
      best = i

      best = l if l < n && [email protected](@heap[best][0], @heap[l][0])
      best = r if r < n && [email protected](@heap[best][0], @heap[r][0])

      break if best == i
      swap(i, best)
      i = best
    end
  end

  def swap(i, j)
    @heap[i], @heap[j] = @heap[j], @heap[i]
  end
end

# Example:
# pq = PriorityQueue.new
# pq.push("B", 1)
# pq.push("C", 3)
# pq.push("A", 5)
# p pq.pop  # => "B"
# p pq.pop  # => "C"
# p pq.pop  # => "A"

# Max-heap variant:
# maxpq = PriorityQueue.new { |a, b| a > b }