|
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 |