Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
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
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
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
You can’t perform that action at this time.