priority_queue.rb

class PriorityQueue

  include Enumerable

  def initialize(from = [], &cmp)
    @cmp = cmp || proc { |a, b| a <=> b }
    @heap = BinaryHeap.new(from, &@cmp)
  end

  [:push, :pop, :empty?].each do |m|
    define_method(m) { |*args| @heap.send(m, *args) }
  end

  def each
    if block_given?
      heap = @heap.clone
      yield heap.pop until heap.empty?
    else
      enum_for(:each)
    end
  end

  def merge(*args)
    @heap.merge(*args)
  end

  alias :concat :merge

  class BinaryHeap

    def initialize(initial_values = [], &cmp)
      @cmp = cmp || proc { |a, b| a <=> b }
      initialize_with(initial_values)
    end

    def push(value)
      @fresh_sorted = false
      @heap.push(value)

      upheap_index = @heap.size - 1
      pi = parent_index(upheap_index)
      while pi && @cmp[@heap[upheap_index], @heap[pi]] == -1
        upheap_index = swap_with_parent(upheap_index)
        pi = parent_index(upheap_index)
      end
    end

    def pop
      # pop the minimum value off
      popped = @heap.shift

      # you would think this unnecessary, except that the next call will
      # push a nil onto the heap if it is empty, making all external calls
      # to empty? return false when in fact they should return true
      return popped if @fresh_sorted || @heap.empty?

      # remove trailing nils from the heap array
      # and move the last value to the root
      @heap.unshift(@heap.pop)

      # run downheap algorithm to restore the heap property and return the
      # originally popped value
      downheap(0)
      popped
    end

    def merge(values)
      @heap += values.to_a
      @heap.sort!
      @fresh_sorted = true
    end

    def empty?
      @heap.empty?
    end

    private

    def initialize_with(values)
      @heap = values.clone.sort
      @fresh_sorted = true
    end

    def downheap(start_index)
      min_child_index = lambda do |pi|
        cis = child_indices(pi)
        cvs = @heap.values_at(*cis).compact
        if cvs.empty?
          nil
        elsif cvs.first == cvs.sort(&@cmp).first
          cis.first
        else
          cis.last
        end
      end

      downheap_index = start_index
      child_index = min_child_index[downheap_index]
      while child_index && @cmp[@heap[downheap_index], @heap[child_index]] == 1
        @heap[child_index], @heap[downheap_index] = 
          @heap[downheap_index], @heap[child_index]
        downheap_index = child_index
        child_index = min_child_index[downheap_index]
      end
    end

    def swap_with_parent(child_index)
      pi = parent_index(child_index)
      @heap[pi], @heap[child_index] = @heap[child_index], @heap[pi]
      pi
    end

    def child_indices(parent_index)
      [2 * parent_index + 1, 2 * parent_index + 2]
    end

    def parent_index(child_index)
      if child_index > 0
        (child_index - 1) / 2
      else
        nil
      end
    end

  end

end

test_priority_queue.rb

require_relative 'priority_queue'
require 'minitest/autorun'

describe PriorityQueue do

  before do
    @pq = PriorityQueue.new
  end

  it "must accept pushed inputs and pop them in the correct order" do
    vals = 1000.times.map { rand(50) }
    pq = PriorityQueue.new(vals)
    # repop some and push them again to be double-sure
    500.times.map { pq.pop }.each { |v| pq.push(v) }
    pq.merge(100.times.map { pq.pop })
    500.times.map { pq.pop }.each { |v| pq.push(v) }
    assert_equal(vals.sort, pq.to_a)
  end

  it "must be faster than PQueue for big chunk operations" do
    require 'pqueue'
    vals = 100000.times.map { rand(1000000) }

    t = Time.now
    pq = PriorityQueue.new(vals)
    pq.pop until pq.empty?
    time = Time.now - t
    
    t = Time.now
    pq = PQueue.new(vals)
    pq.pop until pq.empty?
    pqtime = Time.now - t
    
    puts "PriorityQueue block time: #{time} seconds, PQueue time: #{pqtime} seconds"
    assert(time < pqtime)
  end

  it "must be faster than PQueue for mixed operations" do
    require 'pqueue'
    vals = 100000.times.map { rand(1000000) }

    v = vals.clone
    pq = PriorityQueue.new
    t = Time.now
    until v.empty?
      10.times { pq.push(v.shift) }
      5.times { pq.pop }
    end
    pq.pop until pq.empty?
    time = Time.now - t

    v = vals.clone
    pq = PQueue.new
    t = Time.now
    until v.empty?
      10.times { pq.push(v.shift) }
      5.times { pq.pop }
    end
    pq.pop until pq.empty?
    pqtime = Time.now - t

    puts "PriorityQueue atomic time: #{time} seconds, PQueue time: #{pqtime} seconds"
    assert(time < pqtime)
  end
end